diff --git a/.cargo/config.toml b/.cargo/config.toml index 9ffc1961b09..fd372c95b84 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -7,6 +7,11 @@ linker = "riscv64-linux-gnu-gcc" [alias] sqlness = "run --bin sqlness-runner --target-dir target/sqlness --" +[registry] +# Requires cargo 1.100+; older cargo ignores this key. Only registry crates +# carry a publish time, so git dependencies (our GreptimeTeam forks) are exempt. +global-min-publish-age = "7 days" + [unstable.git] shallow_index = true shallow_deps = true diff --git a/.github/actions/build-linux-artifacts/action.yml b/.github/actions/build-linux-artifacts/action.yml index 3cb9c439551..c5212c28998 100644 --- a/.github/actions/build-linux-artifacts/action.yml +++ b/.github/actions/build-linux-artifacts/action.yml @@ -31,6 +31,22 @@ inputs: description: Build GreptimeDB with large page size (65536). required: false default: 'false' + aws-ci-test-bucket: + description: Bucket used by the object-store integration tests + required: false + default: '' + aws-region: + description: Region of the test bucket + required: false + default: '' + aws-access-key-id: + description: Access key id for the test bucket + required: false + default: '' + aws-secret-access-key: + description: Secret access key for the test bucket + required: false + default: '' runs: using: composite steps: @@ -38,6 +54,13 @@ runs: if: ${{ inputs.disable-run-tests == 'false' }} shell: bash # NOTE: If the BUILD_JOBS > 4, it's always OOM in EC2 instance. + # Without GT_S3_*, every remote object-store case is skipped and the + # release only ever exercises the local file backend. + env: + GT_S3_BUCKET: ${{ inputs.aws-ci-test-bucket }} + GT_S3_REGION: ${{ inputs.aws-region }} + GT_S3_ACCESS_KEY_ID: ${{ inputs.aws-access-key-id }} + GT_S3_ACCESS_KEY: ${{ inputs.aws-secret-access-key }} run: | cd ${{ inputs.working-dir }} && \ make run-it-in-container BUILD_JOBS=4 \ diff --git a/.github/actions/setup-benchmark/action.yml b/.github/actions/setup-benchmark/action.yml new file mode 100644 index 00000000000..a446a78a2c3 --- /dev/null +++ b/.github/actions/setup-benchmark/action.yml @@ -0,0 +1,38 @@ +name: Set up benchmark dependencies +description: Reuse installed tools and install missing dependencies on an Ubuntu self-hosted runner +runs: + using: composite + steps: + - name: Check host and install missing tools + shell: bash + run: | + set -euo pipefail + [[ "$(uname -m)" == x86_64 ]] + command -v git + missing=() + for tool in curl jq python3 rsync; do + command -v "$tool" >/dev/null || missing+=("$tool") + done + if ((${#missing[@]})); then + sudo -n apt-get update + sudo -n env DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l apt-get install -y --no-install-recommends "${missing[@]}" + fi + - name: Install Docker if missing and verify access + shell: bash + run: | + set -euo pipefail + if ! command -v docker >/dev/null; then + . /etc/os-release + [[ "$ID" == ubuntu ]] + sudo -n install -m 0755 -d /etc/apt/keyrings + curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo -n tee /etc/apt/keyrings/docker.asc >/dev/null + sudo -n chmod a+r /etc/apt/keyrings/docker.asc + echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu $VERSION_CODENAME stable" | sudo -n tee /etc/apt/sources.list.d/docker.list >/dev/null + sudo -n apt-get update + sudo -n env DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l apt-get install -y --no-install-recommends docker-ce docker-ce-cli containerd.io docker-buildx-plugin + fi + if ! docker info >/dev/null 2>&1; then + sudo -n systemctl enable --now docker + fi + # Membership must be configured before the runner starts, not in this job. + docker info diff --git a/.github/actions/setup-minio/action.yml b/.github/actions/setup-minio/action.yml index 89fbe28baa7..523203be79f 100644 --- a/.github/actions/setup-minio/action.yml +++ b/.github/actions/setup-minio/action.yml @@ -19,6 +19,8 @@ runs: --set rootUser=rootuser,rootPassword=rootpass123 \ --set buckets[0].name=default \ --set service.port=80,service.targetPort=9000 \ + --set image.repository=pgsty/silo,image.tag=RELEASE.2026-09-16T00-00-00Z \ + --set mcImage.repository=pgsty/mc,mcImage.tag=RELEASE.2026-09-16T00-00-00Z \ minio/minio \ --create-namespace \ -n minio diff --git a/.github/actions/start-runner/action.yml b/.github/actions/start-runner/action.yml index 78568e26d76..44a5ff3feb3 100644 --- a/.github/actions/start-runner/action.yml +++ b/.github/actions/start-runner/action.yml @@ -56,7 +56,7 @@ runs: - name: Start EC2 runner if: startsWith(inputs.runner, 'ec2') - uses: machulav/ec2-github-runner@v2.3.8 + uses: machulav/ec2-github-runner@v2.6.1 id: start-linux-arm64-ec2-runner with: mode: start diff --git a/.github/actions/stop-runner/action.yml b/.github/actions/stop-runner/action.yml index e25a1d8dad7..fb84f2b7181 100644 --- a/.github/actions/stop-runner/action.yml +++ b/.github/actions/stop-runner/action.yml @@ -33,7 +33,7 @@ runs: - name: Stop EC2 runner if: ${{ inputs.label && inputs.ec2-instance-id }} - uses: machulav/ec2-github-runner@v2.3.8 + uses: machulav/ec2-github-runner@v2.6.1 with: mode: stop label: ${{ inputs.label }} diff --git a/.github/runner-scale-sets/query-regression/Dockerfile b/.github/runner-scale-sets/query-regression/Dockerfile index d3e89e81769..617278adf95 100644 --- a/.github/runner-scale-sets/query-regression/Dockerfile +++ b/.github/runner-scale-sets/query-regression/Dockerfile @@ -13,7 +13,7 @@ RUN git init /tmp/otelgen \ -o /usr/local/bin/otelgen ./cmd/otelgen FROM ghcr.io/astral-sh/uv@sha256:3d868e555f8f1dbc324afa005066cd11e1053fc4743b9808ca8025283e65efa5 AS uv -FROM ghcr.io/actions/actions-runner@sha256:08c30b0a7105f64bddfc485d2487a22aa03932a791402393352fdf674bda2c29 +FROM ghcr.io/actions/actions-runner:2.337.0@sha256:e5496277be5d09bc968b3d64911b74e219ac4a3f2edce956a3ecf9271bea1ef4 USER root @@ -22,6 +22,9 @@ ARG SCCACHE_SHA256=aec995a83ad3dff3d14b6314e08858b7b73d35ca85a5bcf3d3a9ec07dee35 ARG RUSTUP_INIT_VERSION=1.29.0 ARG RUSTUP_INIT_TARGET=x86_64-unknown-linux-gnu ARG RUSTUP_INIT_SHA256=4acc9acc76d5079515b46346a485974457b5a79893cfb01112423c89aeb5aa10 +# Pinned to a nightly deliberately: the benchmark compiles historical base +# checkouts (previous nightly builds) that still use #![feature] gates, so +# the runner cannot follow the workspace's rust-toolchain.toml (now stable). ARG RUST_TOOLCHAIN=nightly-2026-03-21 ENV RUSTUP_HOME=/opt/rustup \ diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index fa8d436d2c2..44500a8e50f 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -86,15 +86,35 @@ sweep. ### Building and updating the ECS image The runner `Dockerfile` in the parent directory stays the single source of the -tool contract. Build a new ECS image from it: +tool contract. The image is rebuilt **automatically** by the +`rebuild-query-regression-runner-image` job in +`.github/workflows/release-dev-builder-images.yaml` whenever anything under +this directory changes on main (or via manual dispatch): it runs the ops tool +below, opens an epoch-bump PR for `RUNNER_IMAGE_EPOCH` in +`query-regression.yml`, and points the `QUERY_REGRESSION_ECS_IMAGE_ID` repo +variable at the new image. + +The runner toolchain is **pinned inside the Dockerfile** to +`nightly-2026-03-21` and deliberately does NOT follow the workspace +`rust-toolchain.toml`: the benchmark also compiles the BASE checkout (the +previous nightly build), which may predate the stable-toolchain migration and +still require a nightly compiler. + +The manual fallback (also what the workflow runs): ```bash ALIBABA_CLOUD_ACCESS_KEY_ID=... ALIBABA_CLOUD_ACCESS_KEY_SECRET=... \ uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ --region-id --vswitch-id --security-group-id \ - --base-image-id + [--base-image-id ] ``` +`--base-image-id` is optional: the script defaults to the latest public +Ubuntu 26.04 image in the region (26.04 matches the tool-version pins the +Verify step asserts, e.g. python3 3.14; keep the two in sync); pass it — or +set the `ALIYUN_ECS_BASE_IMAGE_ID` repo variable consumed by the automated +job — to pin a specific base image. + The script boots a temporary builder instance, `docker build`s the runner image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and `/home/runner` (actions-runner) onto the host, installs the ephemeral-runner @@ -158,10 +178,16 @@ overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`) and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`, `Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and -the image-baked `nightly-2026-03-21` Rust toolchain. `mold` and `python3` -come from apt at image-build time (not Ubuntu 24.04's default 3.12); bump -the Verify pins together with `QUERY_REGRESSION_ECS_IMAGE_ID` when the -image is rebuilt. Rustup, Cargo, and Rustc +the image-baked Rust toolchain `nightly-2026-03-21` pinned in the runner +Dockerfile (deliberately independent of `rust-toolchain.toml` — see above). +`mold` and `python3` +come from apt at image-build time; if the +Ubuntu archive ships a newer package revision between rebuilds, the Verify +step fails with the observed version — bump those pins in `query-regression.yml` +when that happens. Everything else (toolchain, uv, sccache, otelgen, rustup, +the runner base) is pinned by digest/sha/commit in the Dockerfile, and the +image id + `RUNNER_IMAGE_EPOCH` updates are automated by the rebuild job, so +a routine rebuild needs no manual file updates. Rustup, Cargo, and Rustc must resolve from `/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. Protobuf well-known includes, including `google/protobuf/any.proto` and `google/protobuf/empty.proto`, are an image diff --git a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py index 7031917f910..f9ecb6b2333 100644 --- a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py +++ b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py @@ -25,7 +25,7 @@ """Build the query-regression ECS custom image (manual ops tool). -Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 24.04 image, +Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 26.04 image, builds the existing runner container image (the Dockerfile in the parent directory remains the single source of the tool contract), materializes the tool directories onto the host filesystem so the workflow's "Verify runner @@ -58,9 +58,9 @@ POLL_INTERVAL_SECONDS = 15 CONSOLE_POLL_INTERVAL_SECONDS = 30 BUILD_TIMEOUT_SECONDS = 60 * 60 -# Same apt package contract as the runner Dockerfile; the base -# actions-runner image is Ubuntu 24.04, so an Ubuntu 24.04 host resolves the -# same tool versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). +# Same apt package contract as the runner Dockerfile; the base is an +# Ubuntu 26.04 image, so an Ubuntu 26.04 host resolves the same tool +# versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). # Docker itself comes from Docker's official repository (docker-ce), not the # distribution-packaged docker.io. APT_PACKAGES = [ @@ -229,6 +229,64 @@ def call_api_with_retry(fn, description: str, attempts: int = 5): raise RuntimeError("unreachable: retry loop exited without returning") +def resolve_base_image_id(client, region_id: str) -> str: + """Resolve the latest public Ubuntu 26.04 x86_64 system image. + + Used as the default for --base-image-id: the runner Dockerfile pins + every tool version itself, so a current stock Ubuntu LTS base is all + the builder needs. Pass --base-image-id (or ALIYUN_ECS_BASE_IMAGE_ID) + to pin a specific base image deterministically. + """ + from alibabacloud_ecs20140526 import models as ecs_models + + def _is_target_ubuntu(image) -> bool: + # osname is localized (e.g. "Ubuntu 26.04 64位"), osname_en the + # English form; accept either. Keep the Ubuntu version in sync with + # the tool-version pins asserted by the Verify step in + # query-regression.yml (e.g. python3 3.14 comes from 26.04). + for os_name in (image.osname_en, image.osname): + if os_name and "Ubuntu" in os_name and "26.04" in os_name: + return True + return False + + page_size = 100 + images: list = [] + page_number = 1 + while True: + page = call_api_with_retry( + lambda: client.describe_images( + ecs_models.DescribeImagesRequest( + region_id=region_id, + image_owner_alias="system", + ostype="linux", + architecture="x86_64", + page_number=page_number, + page_size=page_size, + ) + ), + f"DescribeImages(system base, page {page_number})", + ).body.images.image or [] + images.extend(page) + if len(page) < page_size: + break + page_number += 1 + + candidates = [image for image in images if _is_target_ubuntu(image)] + if not candidates: + raise SystemExit( + "No public Ubuntu 26.04 x86_64 system image found in region " + f"{region_id}; pass --base-image-id explicitly" + ) + candidates.sort(key=lambda image: image.creation_time or "", reverse=True) + picked = candidates[0] + print( + f"Resolved base image: {picked.image_id} ({picked.osname_en or picked.osname}, " + f"created {picked.creation_time}) from {len(candidates)} candidates", + flush=True, + ) + return picked.image_id + + def read_console_output(client, region_id: str, instance_id: str) -> str: """Fetch the instance's serial console output; no in-guest agent needed.""" from alibabacloud_ecs20140526 import models as ecs_models @@ -253,13 +311,19 @@ def main() -> int: parser.add_argument("--image-name", default=None, help="Defaults to a timestamped name.") args = parser.parse_args() - for name in ("region_id", "vswitch_id", "security_group_id", "base_image_id"): + for name in ("region_id", "vswitch_id", "security_group_id"): if not getattr(args, name): raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") from alibabacloud_ecs20140526 import models as ecs_models client = make_ecs_client(args.region_id) + # --base-image-id is optional: default to the latest public Ubuntu 26.04 + # image in the region (the Dockerfile pins every tool version itself, so + # base drift is low-risk; pass --base-image-id or set + # ALIYUN_ECS_BASE_IMAGE_ID to pin deterministically). + if not args.base_image_id: + args.base_image_id = resolve_base_image_id(client, args.region_id) image_name = args.image_name or time.strftime( "greptimedb-query-regression-runner-%Y%m%d%H%M%S", time.gmtime() ) diff --git a/.github/scripts/agent-observability-summary.py b/.github/scripts/agent-observability-summary.py index 38ba4e9dbdd..33ba3ad7e6e 100644 --- a/.github/scripts/agent-observability-summary.py +++ b/.github/scripts/agent-observability-summary.py @@ -18,6 +18,8 @@ import json import math from pathlib import Path +from benchmark_resource_summary import render_resources + TARGETS = ('greptimedb', 'clickhouse', 'victorialogs') @@ -201,6 +203,7 @@ def lifecycle(root, targets): rows.append([context, entry.get('phase'), duration(entry.get('elapsed_ms')) if valid else 'invalid timestamp units (legacy artifact)', entry.get('exit_code')]) lines += ['## Phase timings', '', 'Phase duration includes orchestration overhead. Container cleanup does not prove ECS teardown.', '', table(['DB / scope', 'Phase', 'Duration', 'Exit code'], rows), ''] + lines.append(render_resources([('dataset', root)] + [(t, root / f'{t}-1') for t in targets], table)) return '\n'.join(lines) diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py index 1a26776b4cd..46efbffd3ec 100644 --- a/.github/scripts/aliyun-ecs-runner-provision.py +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -49,10 +49,10 @@ import base64 import json import os import time -import urllib.error -import urllib.request from dataclasses import dataclass +from runner_utils import append_github_output, create_registration_token, find_runner_by_name, github_api + RUNNER_NAME_PREFIX = "qreg-ecs" RUNNER_LABEL_PREFIX = "query-regression-ecs" MANAGED_BY_TAG_KEY = "managed-by" @@ -243,63 +243,6 @@ def encode_user_data(script: str) -> str: return base64.b64encode(script.encode("utf-8")).decode("ascii") -def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: - data = json.dumps(body).encode("utf-8") if body is not None else None - request = urllib.request.Request( - f"https://api.github.com{path}", - data=data, - method=method, - headers={ - "Accept": "application/vnd.github+json", - "Authorization": f"Bearer {token}", - "X-GitHub-Api-Version": "2022-11-28", - }, - ) - try: - with urllib.request.urlopen(request, timeout=30) as response: - if response.status == 204: - return {} - return json.loads(response.read().decode("utf-8")) - except urllib.error.HTTPError as error: - # GitHub's error body says exactly why (e.g. "Must have admin rights to - # Repository" for a PAT without the required scope); surface it instead - # of a bare "HTTP Error 403". - body = error.read().decode("utf-8", "replace") - raise SystemExit( - f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" - "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " - "'repo' scope (classic PAT) or 'Administration: write' on the " - "repository (fine-grained PAT)." - ) from error - - -def create_registration_token(github_token: str, repo: str) -> str: - response = github_api( - github_token, - "POST", - f"/repos/{repo}/actions/runners/registration-token", - body={}, - ) - return response["token"] - - -def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: - page = 1 - while True: - response = github_api( - github_token, - "GET", - f"/repos/{repo}/actions/runners?per_page=100&page={page}", - ) - runners = response.get("runners", []) - for runner in runners: - if runner.get("name") == name: - return runner - if len(runners) < 100: - return None - page += 1 - - def make_ecs_client(config: ProvisionConfig): from alibabacloud_ecs20140526.client import Client as EcsClient from alibabacloud_tea_openapi.models import Config as OpenApiConfig @@ -388,13 +331,6 @@ def dump_console_output(client, region_id: str, instance_id: str) -> None: summary.write("\n```\n\n") -def append_github_output(name: str, value: str) -> None: - output_path = os.environ.get("GITHUB_OUTPUT") - if output_path: - with open(output_path, "a", encoding="utf-8") as output: - output.write(f"{name}={value}\n") - - def run_instance(client, config: ProvisionConfig, user_data: str) -> str: from alibabacloud_ecs20140526 import models as ecs_models diff --git a/.github/scripts/aliyun-ecs-runner-teardown.py b/.github/scripts/aliyun-ecs-runner-teardown.py index 658633e0395..d5b30b2922d 100644 --- a/.github/scripts/aliyun-ecs-runner-teardown.py +++ b/.github/scripts/aliyun-ecs-runner-teardown.py @@ -47,6 +47,8 @@ import sys from datetime import datetime, timedelta, timezone from pathlib import Path +from runner_utils import deregister_runner + # Reuse the GitHub API client, runner lookup, and tag constants from the # provision script. _PROVISION_SPEC = importlib.util.spec_from_file_location( @@ -159,20 +161,6 @@ def delete_instance(client, instance_id: str, region_id: str | None = None) -> b return False -def deregister_runner(token: str, repo: str, runner_name: str) -> bool: - runner = provision.find_runner_by_name(token, repo, runner_name) - if runner is None: - print(f"Runner {runner_name} is not registered", flush=True) - return True - try: - provision.github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") - print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) - return True - except Exception as error: # noqa: BLE001 - print(f"Failed to deregister runner {runner_name}: {error}", flush=True) - return False - - def list_managed_instances( client, region_id: str ) -> list[tuple[str, str, str, str | None]]: diff --git a/.github/scripts/aws-ec2-runner-provision.py b/.github/scripts/aws-ec2-runner-provision.py new file mode 100644 index 00000000000..7bf1b4091c9 --- /dev/null +++ b/.github/scripts/aws-ec2-runner-provision.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# /// script +# requires-python = ">=3.10" +# dependencies = ["boto3>=1.35,<2"] +# /// + +"""Provision an ephemeral EC2 CI runner from an Ubuntu x86_64 AMI.""" + +from __future__ import annotations + +import argparse +import os +import re +import shlex +import time +from dataclasses import dataclass + +from runner_utils import append_github_output, create_registration_token, find_runner_by_name, github_api + +MANAGED_BY = "query-regression-ci" +REPO_TAG = "github-repository" +RUN_TAG = "query-regression-run-id" +RUNNER_TAG = "github-runner-name" +TTL_TAG_KEY = "runner-ttl-hours" + + +def make_client(region: str): + import boto3 + from botocore.config import Config + + return boto3.client("ec2", region_name=region, config=Config( + connect_timeout=10, read_timeout=30, + retries={"mode": "standard", "max_attempts": 3}, + )) + + +@dataclass(frozen=True) +class ProvisionConfig: + region: str + image_id: str + instance_type: str + subnet_id: str + security_group_id: str + repo: str + run_id: str + github_token: str + system_disk_gib: int = 500 + ttl_hours: int = 8 + + def __post_init__(self): + if not re.fullmatch(r"[0-9]+-[0-9]+", self.run_id): + raise ValueError("run-id must be GitHub run_id-run_attempt") + if not 20 <= self.system_disk_gib <= 2048: + raise ValueError("system-disk-gib must be between 20 and 2048") + if not 1 <= self.ttl_hours <= 168: + raise ValueError("ttl-hours must be between 1 and 168") + + @property + def runner_name(self) -> str: + return f"ci-ec2-{self.run_id}" + + +def resolve_image(client, image_id: str) -> str: + if image_id != "auto": + return image_id + images = client.describe_images( + Owners=["099720109477"], + Filters=[ + {"Name": "name", "Values": ["ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*"]}, + {"Name": "state", "Values": ["available"]}, + {"Name": "architecture", "Values": ["x86_64"]}, + ], + )["Images"] + if not images: + raise ValueError("No Canonical Ubuntu 24.04 amd64 AMI found in this region") + image = max(images, key=lambda image: image["CreationDate"]) + print(f"Using AMI {image['ImageId']}: {image['Name']}") + return image["ImageId"] + + +def root_disk(client, image_id: str, size: int) -> dict: + images = client.describe_images(ImageIds=[image_id])["Images"] + if len(images) != 1: + raise ValueError(f"AMI not found: {image_id}") + image = images[0] + if image["State"] != "available" or image["Architecture"] != "x86_64": + raise ValueError("Runner AMI must be available and x86_64") + root = next((d for d in image["BlockDeviceMappings"] + if d["DeviceName"] == image["RootDeviceName"] and "Ebs" in d), None) + if root is None: + raise ValueError("Runner AMI must have an EBS root disk") + minimum = root["Ebs"]["VolumeSize"] + if size < minimum: + raise ValueError(f"Requested root disk {size} GiB is smaller than AMI minimum {minimum} GiB") + return {"DeviceName": image["RootDeviceName"], "Ebs": { + "VolumeSize": size, "VolumeType": "gp3", "Iops": 3000, + "Throughput": 125, "DeleteOnTermination": True, + }} + + +def runner_download(token: str, repo: str) -> str: + runners = github_api(token, "GET", f"/repos/{repo}/actions/runners/downloads") + return next(r["download_url"] for r in runners if r["os"] == "linux" and r["architecture"] == "x64") + + +def render_user_data(name: str, token: str, repo: str, download_url: str) -> str: + # Quote every external value before embedding it in the root shell script. + name, token, repo, download_url = map(shlex.quote, (name, token, repo, download_url)) + return f"""#!/bin/bash +set -euo pipefail +for tool in git curl tar sudo; do command -v "$tool"; done +# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole +# service cgroup on PSI pressure, and swap thrashing looks like pressure. +systemctl disable --now systemd-oomd.socket systemd-oomd.service || true +systemctl mask systemd-oomd.socket systemd-oomd.service || true +# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services +# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub +# records UserCancelled (job still valid, child exit 143). These VMs live +# ~1h; security updates belong in the image, not at job time. +systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl stop unattended-upgrades.service || true +killall -9 unattended-upgr || true +cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF' +APT::Periodic::Update-Package-Lists "0"; +APT::Periodic::Unattended-Upgrade "0"; +APT::Periodic::Download-Upgradeable-Packages "0"; +APTEOF +if [[ ! -f "/swapfile" ]]; then + fallocate --length 16G "/swapfile" + chmod 600 "/swapfile" + mkswap "/swapfile" +fi +swapon "/swapfile" +sysctl --write vm.swappiness=10 +swapon --show +free --human + +id runner >/dev/null 2>&1 || useradd --create-home --shell /bin/bash runner +# Establish membership before the runner service starts; setup installs Docker later. +getent group docker >/dev/null || groupadd --system docker +usermod -aG docker runner +printf 'runner ALL=(ALL) NOPASSWD: ALL\\n' > /etc/sudoers.d/ci-runner +chmod 0440 /etc/sudoers.d/ci-runner +visudo -cf /etc/sudoers.d/ci-runner +install -d -o runner -g "$(id -gn runner)" /home/runner/benchmark-data /opt/ci-github-runner +cd /opt/ci-github-runner +curl --fail --location --retry 3 {download_url} | tar xz +./bin/installdependencies.sh +chown -R runner:"$(id -gn runner)" /opt/ci-github-runner +runuser -u runner -- ./config.sh --unattended --ephemeral --disableupdate --url https://github.com/{repo} --token {token} --name {name} --labels {name} +./svc.sh install runner +# Match the ECS runner's behavior when a child process is OOM-killed. +service_name=$(cat .service) +mkdir -p "/etc/systemd/system/$service_name.d" +cat > "/etc/systemd/system/$service_name.d/oom.conf" <<'EOF' +[Service] +OOMPolicy=continue +EOF +systemctl daemon-reload +# Package installation in a job must not restart its runner service. +install -d /etc/needrestart/conf.d +cat > /etc/needrestart/conf.d/actions_runner_services.conf <<'EOF' +$nrconf{{override_rc}}{{qr(^actions\\.runner\\..+\\.service$)}} = 0; +EOF +./svc.sh start +""" + + +def run_instance(client, config: ProvisionConfig) -> str: + image_id = resolve_image(client, config.image_id) + disk = root_disk(client, image_id, config.system_disk_gib) + instance = client.describe_instance_types(InstanceTypes=[config.instance_type])["InstanceTypes"][0] + if "x86_64" not in instance["ProcessorInfo"]["SupportedArchitectures"]: + raise ValueError("Runner instance type must support x86_64") + token = create_registration_token(config.github_token, config.repo) + print(f"::add-mask::{token}") + script = render_user_data(config.runner_name, token, config.repo, runner_download(config.github_token, config.repo)) + tags = [ + {"Key": "Name", "Value": config.runner_name}, + {"Key": "managed-by", "Value": MANAGED_BY}, + {"Key": REPO_TAG, "Value": config.repo}, + {"Key": RUN_TAG, "Value": config.run_id}, + {"Key": RUNNER_TAG, "Value": config.runner_name}, + {"Key": TTL_TAG_KEY, "Value": str(config.ttl_hours)}, + ] + response = client.run_instances( + ImageId=image_id, InstanceType=config.instance_type, + MinCount=1, MaxCount=1, SubnetId=config.subnet_id, + SecurityGroupIds=[config.security_group_id], + BlockDeviceMappings=[disk], UserData=script, + ClientToken=config.runner_name, + MetadataOptions={"HttpTokens": "required", "HttpEndpoint": "enabled"}, + TagSpecifications=[{"ResourceType": kind, "Tags": tags} for kind in ("instance", "volume")], + ) + return response["Instances"][0]["InstanceId"] + + +def dump_console(client, instance_id: str) -> None: + try: + # Boto3 decodes the API's base64 console output. + output = client.get_console_output(InstanceId=instance_id, Latest=True).get("Output", "") + print(f"::group::EC2 console {instance_id}") + print("\n".join(output.splitlines()[-80:])) + print("::endgroup::") + except Exception as error: + print(f"Unable to read console for {instance_id}: {error}") + + +def provision(client, config: ProvisionConfig) -> int: + instance_id = run_instance(client, config) + # Publish before waiting: failed registration must still reach teardown. + for key, value in (("instance_id", instance_id), + ("label", config.runner_name), ("runner_name", config.runner_name)): + append_github_output(key, value) + try: + client.get_waiter("instance_running").wait( + InstanceIds=[instance_id], WaiterConfig={"Delay": 5, "MaxAttempts": 60}) + deadline = time.monotonic() + 10 * 60 + while time.monotonic() < deadline: + runner = find_runner_by_name(config.github_token, config.repo, config.runner_name) + if runner and runner.get("status") == "online": + print(f"Runner {config.runner_name} online: {instance_id}") + return 0 + time.sleep(5) + raise TimeoutError(f"Runner {config.runner_name} did not come online") + finally: + dump_console(client, instance_id) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + for name, env in (("region", "AWS_REGION"), ("repo", "GITHUB_REPOSITORY"), + ("github-token", "GH_PERSONAL_ACCESS_TOKEN"), + ("run-id", "AWS_EC2_RUN_ID"), ("image-id", "AWS_EC2_IMAGE_ID"), + ("instance-type", "AWS_EC2_INSTANCE_TYPE"), + ("subnet-id", "AWS_EC2_SUBNET_ID"), + ("security-group-id", "AWS_EC2_SECURITY_GROUP_ID")): + parser.add_argument(f"--{name}", default=os.environ.get(env) or ("auto" if name == "image-id" else None)) + parser.add_argument("--system-disk-gib", type=int, default=os.environ.get("AWS_EC2_SYSTEM_DISK_GIB", "500")) + parser.add_argument("--ttl-hours", type=int, default=os.environ.get("AWS_EC2_TTL_HOURS", "8")) + args = parser.parse_args() + for name, value in vars(args).items(): + if value is None or value == "": + parser.error(f"Missing {name.replace('_', '-')}") + config = ProvisionConfig(**vars(args)) + return provision(make_client(args.region), config) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/aws-ec2-runner-teardown.py b/.github/scripts/aws-ec2-runner-teardown.py new file mode 100644 index 00000000000..fbfba4f639a --- /dev/null +++ b/.github/scripts/aws-ec2-runner-teardown.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# /// script +# requires-python = ">=3.10" +# dependencies = ["boto3>=1.35,<2"] +# /// + +"""Delete managed EC2 CI runners, or sweep expired instances by ownership tags.""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +import time +from pathlib import Path +from datetime import datetime, timedelta, timezone + +from runner_utils import deregister_runner + +_SPEC = importlib.util.spec_from_file_location("aws_ec2_runner_provision", Path(__file__).with_name("aws-ec2-runner-provision.py")) +assert _SPEC is not None and _SPEC.loader is not None +provision = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = provision +_SPEC.loader.exec_module(provision) +make_client = provision.make_client +dump_console = provision.dump_console +MANAGED_BY = provision.MANAGED_BY +REPO_TAG = provision.REPO_TAG +RUN_TAG = provision.RUN_TAG +RUNNER_TAG = provision.RUNNER_TAG +TTL_TAG_KEY = provision.TTL_TAG_KEY +SWEEP_TTL_HOURS = 4 + + +def managed_instances(client, repo: str, run_id: str | None = None, + instance_id: str | None = None) -> list[dict]: + filters = [{"Name": "tag:managed-by", "Values": [MANAGED_BY]}, + {"Name": f"tag:{REPO_TAG}", "Values": [repo]}] + if instance_id: + # A retried teardown may have a newer workflow attempt than creation. + filters.append({"Name": "instance-id", "Values": [instance_id]}) + elif run_id: + filters.append({"Name": f"tag:{RUN_TAG}", "Values": [run_id]}) + return [instance + for page in client.get_paginator("describe_instances").paginate(Filters=filters) + for reservation in page["Reservations"] for instance in reservation["Instances"]] + + +def delete_instance(client, instance_id: str) -> bool: + dump_console(client, instance_id) + for attempt in range(12): + try: + client.terminate_instances(InstanceIds=[instance_id]) + client.get_waiter("instance_terminated").wait( + InstanceIds=[instance_id], WaiterConfig={"Delay": 5, "MaxAttempts": 60}) + return True + except Exception as error: + code = getattr(error, "response", {}).get("Error", {}).get("Code") + if code == "InvalidInstanceID.NotFound": + if attempt < 11: + time.sleep(5) + continue + # The ownership lookup observed this instance. Treat sustained + # NotFound after bounded retries as an already-deleted VM. + return True + print(f"Failed to terminate {instance_id}: {error}") + return False + return False + + +def teardown(client, repo: str, token: str, run_id: str | None = None, + instance_id: str | None = None, now: datetime | None = None, + runner_name: str | None = None) -> int: + # Lookup by ownership/run tags also recovers a create whose response or job + # outputs were lost. Sweep never touches legacy build/JSONBench instances. + instances = managed_instances(client, repo, run_id, instance_id) + # Tags can lag RunInstances. Retry before concluding a targeted VM is gone. + if run_id and not instances: + for _ in range(11): + time.sleep(5) + instances = managed_instances(client, repo, run_id, instance_id) + if instances: + break + ok = True + found = False + for instance in instances: + identity = instance["InstanceId"] + if instance_id and identity != instance_id: + continue + found = True + tags = {t["Key"]: t["Value"] for t in instance.get("Tags", [])} + name = tags.get(RUNNER_TAG, "") + if not name.startswith("ci-ec2-"): + print(f"Skipping {identity}: missing managed runner name") + ok = False + continue + if not run_id: + ttl = tags.get(TTL_TAG_KEY) + try: + hours = SWEEP_TTL_HOURS if ttl is None else int(ttl) + if not 1 <= hours <= 168: + raise ValueError("TTL outside 1..168 hours") + except (TypeError, ValueError): + print(f"Skipping {identity}: invalid TTL {ttl!r}") + continue + if (now or datetime.now(timezone.utc)) - instance["LaunchTime"] < timedelta(hours=hours): + continue + # A GitHub outage must not prevent EC2 deletion, or vice versa. + ok &= delete_instance(client, identity) + try: + ok &= deregister_runner(token, repo, name) + except (Exception, SystemExit) as error: + print(f"Failed to unregister {name}: {error}") + ok = False + if run_id and not found: + # EC2 may already be gone while its GitHub registration remains. + name = runner_name or f"ci-ec2-{run_id}" + if not name.startswith("ci-ec2-"): + raise ValueError("Invalid CI runner name") + ok &= deregister_runner(token, repo, name) + return 0 if ok else 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + for name, env in (("region", "AWS_REGION"), ("repo", "GITHUB_REPOSITORY"), + ("github-token", "GH_PERSONAL_ACCESS_TOKEN"), + ("run-id", "AWS_EC2_RUN_ID"), ("instance-id", "AWS_EC2_INSTANCE_ID"), + ("runner-name", "AWS_EC2_RUNNER_NAME")): + parser.add_argument(f"--{name}", default=os.environ.get(env)) + parser.add_argument("--sweep", action="store_true") + args = parser.parse_args() + for name in ("region", "repo", "github_token"): + if not getattr(args, name): + parser.error(f"Missing {name.replace('_', '-')}") + if not args.sweep and not args.run_id: + parser.error("Missing run-id for targeted teardown") + return teardown(make_client(args.region), args.repo, args.github_token, + None if args.sweep else args.run_id, + None if args.sweep else args.instance_id, + runner_name=None if args.sweep else args.runner_name) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/benchmark-version-comparison.py b/.github/scripts/benchmark-version-comparison.py new file mode 100644 index 00000000000..1aa744df9ad --- /dev/null +++ b/.github/scripts/benchmark-version-comparison.py @@ -0,0 +1,191 @@ +# 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. + +"""Compare validated GreptimeDB query results; never gate CI on noisy latency deltas.""" +import argparse +import html +import json +import math +import os +from pathlib import Path +import re +import statistics + + +def read(path): + return json.loads(path.read_text()) if path.is_file() else {} + + +def validate(env): + enabled = env.get('COMPARE_GREPTIMEDB', 'false') + if enabled not in ('true', 'false'): + raise ValueError('COMPARE_GREPTIMEDB must be true or false') + if enabled == 'false': + return + if env.get('RUN_GREPTIMEDB') != 'true': + raise ValueError('Version comparison requires Run GreptimeDB') + for key in ('BASELINE_GREPTIMEDB_TAG', 'GREPTIMEDB_TAG'): + if not re.fullmatch(r'[a-zA-Z0-9_][a-zA-Z0-9_.-]{0,127}(@sha256:[0-9a-f]{64})?', env.get(key, '')): + raise ValueError(f'Invalid {key}') + if env['BASELINE_GREPTIMEDB_TAG'] == env['GREPTIMEDB_TAG']: + raise ValueError('Baseline and candidate must have different image tags') + + +def same(a, b, fields): + return all(a.get(k) is not None and a.get(k) == b.get(k) for k in fields) + + +def finite(value): + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) and value >= 0 + + +def improvement(old, new, throughput=False): + if not finite(old) or not finite(new) or old <= 0: + return None + return (new / old - 1 if throughput else 1 - new / old) * 100 + + +def row(query, metric, old, new, compatible, reason): + delta = improvement(old, new, metric == 'QPS') if compatible else None + return dict(query=query, metric=metric, baseline=old, candidate=new, + improvement_percent=delta, status='comparable' if delta is not None else reason) + + +def agent(root): + def load(r): + return read(r / 'greptimedb-1/greptimedb-1/measured-summary.json') + a, b = load(root / 'baseline'), load(root) + finished = [r / 'greptimedb-1/exit-code.txt' for r in (root / 'baseline', root)] + compatible = (all(p.is_file() and p.read_text().strip() == '0' for p in finished) + and a.get('passed') is True and b.get('passed') is True + and same(a, b, ('contract_sha256', 'data_profile', 'query_window', 'warmup_runs', 'measured_runs'))) + qa, qb = ({q['query_id']: q for q in x.get('per_query', [])} for x in (a, b)) + rows = [] + for q in sorted(qa.keys() | qb.keys()): + old, new = qa.get(q, {}), qb.get(q, {}) + ok = compatible and same(old, new, ('result_fingerprint', 'success')) and old.get('success', 0) > 0 + ok = ok and all(v.get('errors') == 0 and v.get('timeouts') == 0 for v in (old, new)) + for key in ('p50_ms', 'p95_ms', 'p99_ms'): + rows.append(row(q, key, old.get(key), new.get(key), ok, 'failed, missing or incompatible result')) + return rows + + +def traces(root): + a, b = (read(r / 'runs/greptimedb/run.json') for r in (root / 'baseline', root)) + compatible = (a.get('success') is True and b.get('success') is True + and same(a, b, ('expected', 'protocol', 'audit_expected', 'corpus')) + and same(a.get('execution', {}), b.get('execution', {}), + ('db_cpus', 'db_memory', 'load_workers_requested', 'max_in_flight'))) + rows = [] + for q in (f'T{i:02}' for i in range(1, 9)): + values, identities, counts = [], [], [] + ok = compatible + for run in (a, b): + correctness = [s for s in run.get('correctness', []) if s.get('query_id') == q] + all_samples = [s for s in run.get('samples', []) if s.get('query_id') == q] + measured = [s for s in all_samples if s.get('phase') == 'measured'] + passed = bool(correctness and measured) and all(s.get('outcome') == 'passed' for s in correctness + all_samples) + expected = run.get('expected', {}).get('contracts', {}).get(q, {}) + passed = passed and all(s.get('fingerprint') == expected.get('fingerprint') and s.get('fingerprint') is not None + and s.get('row_count') == expected.get('row_count') for s in correctness + all_samples) + elapsed = [s.get('elapsed_ms') for s in measured] + passed = passed and all(finite(v) for v in elapsed) + ok = ok and passed + values.append(statistics.median(elapsed) if passed else None) + identities.append(expected) + counts.append((len(correctness), len(all_samples), len(measured))) + ok = ok and identities[0] == identities[1] and counts[0] == counts[1] + rows.append(row(q, 'median_ms', *values, ok, 'failed, skipped, missing or incompatible result')) + return rows + + +def long_range(root): + rows = [] + manifest = read(root / 'run-manifest.json') + healthy = all( + all(field in (manifest.get(phase, {}).get(target, {}).get('docker', {}) if phase == 'post_load_results' + else manifest.get(phase, {}).get(target, {})).get('state', '').split() + for field in ('Running=true', 'ExitCode=0', 'OOMKilled=false')) + for phase in ('post_load_results', 'post_query_results') + for target in ('greptimedb-baseline', 'greptimedb')) + for mode in ('serial', 'concurrent-duration-c1', 'concurrent-duration-c4'): + a, b = (read(root / name / 'results-vmbench-long-range' / (mode + '-summary.json')) + for name in ('greptimedb-baseline', 'greptimedb')) + fields = ('workload', 'query_file_sha256', 'start', 'end', 'step', 'profile', 'data_duration', + 'scrape_interval', 'cache_policy_effective', 'mode', 'connection_mode', 'request_timeout_seconds') + fields += ('tries', 'warmup_runs') if mode == 'serial' else ('concurrency', 'duration_seconds', 'warmup_seconds') + compatible = healthy and same(a, b, fields) and all(x.get('errors') == 0 and x.get('timeouts') == 0 for x in (a, b)) + qa, qb = ({q['query_index']: q for q in x.get('per_query', [])} for x in (a, b)) + for q in sorted(qa.keys() | qb.keys()): + old, new = qa.get(q, {}), qb.get(q, {}) + ok = compatible and same(old, new, ('query', 'series', 'points')) + ok = ok and all(v.get('success', 0) > 0 and v.get('series', 0) > 0 and v.get('points', 0) > 0 + and v.get('errors') == 0 and v.get('timeouts') == 0 for v in (old, new)) + for key in ('p50_ms', 'qps') if mode != 'serial' else ('p50_ms',): + rows.append(row(f'{mode}/Q{q}', 'QPS' if key == 'qps' else key, + old.get(key), new.get(key), ok, 'failed, missing or incompatible result')) + return rows + + +def render(document): + def cell(value): + return html.escape(str(value)).replace('|', '|').replace('\n', ' ') + def value(v, metric): + if not finite(v): + return '—' + if metric == 'QPS': + return f'{v:.3f}' + return f'{v / 1000:.2f} s' if v >= 1000 else f'{v:.2f} ms' + lines = ['# GreptimeDB version comparison', '', + f'Baseline: `{cell(document["baseline_tag"])}` → candidate: `{cell(document["candidate_tag"])}`.', '', + 'Same-host sequential runs with fresh databases and shared input data. Positive = improvement; ' + 'latency: (baseline − candidate) / baseline; QPS: (candidate − baseline) / baseline. ' + 'Single runs are diagnostic, not a statistically established regression. No performance threshold gates CI.', '', + '| Query | Metric | Baseline | Candidate | Improvement | Status |', + '| --- | --- | --- | --- | --- | --- |'] + for r in document['rows']: + delta = r['improvement_percent'] + lines.append('| ' + ' | '.join(map(cell, [r['query'], r['metric'], value(r['baseline'], r['metric']), + value(r['candidate'], r['metric']), f'{delta:+.1f}%' if delta is not None else '—', r['status']])) + ' |') + if document.get('kind') == 'long-range': + lines += ['', 'Long-range retains the existing non-empty series/point-count checks; these do not establish equality of every returned value.'] + if not document['rows']: + lines += ['', 'Results unavailable; comparison not performed.'] + return '\n'.join(lines) + '\n' + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('command', choices=('validate', 'report')) + parser.add_argument('--kind', choices=('agent', 'traces', 'long-range')) + parser.add_argument('--root', type=Path) + args = parser.parse_args() + if args.command == 'validate': + validate(os.environ) + return + if args.root is None or args.kind is None: + parser.error('report requires --root and --kind') + document = {'kind': args.kind, 'baseline_tag': os.environ.get('BASELINE_GREPTIMEDB_TAG', ''), + 'candidate_tag': os.environ.get('GREPTIMEDB_TAG', ''), + 'rows': {'agent': agent, 'traces': traces, 'long-range': long_range}[args.kind](args.root)} + text = render(document) + # Preserve partial comparisons even if a target failed before creating output. + args.root.mkdir(parents=True, exist_ok=True) + (args.root / 'version-comparison.json').write_text(json.dumps(document, indent=2, allow_nan=False) + '\n') + (args.root / 'version-comparison.md').write_text(text) + print(text, end='') + + +if __name__ == '__main__': + main() diff --git a/.github/scripts/benchmark_resource_summary.py b/.github/scripts/benchmark_resource_summary.py new file mode 100644 index 00000000000..3a0a896c46e --- /dev/null +++ b/.github/scripts/benchmark_resource_summary.py @@ -0,0 +1,99 @@ +# 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. + +"""Render shared O11yBench sampled resource artifacts for lifecycle summaries.""" +import json +import math + + +def number(value): + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) and value >= 0 + + +def metric(values, key, statistic): + entry = values.get(key, {}) + value = entry.get(statistic) + return value if number(entry.get('samples')) and entry['samples'] > 0 and number(value) else None + + +def decimal(value, suffix=''): + return f'{value:.2f}{suffix}' if number(value) else '—' + + +def memory(value): + if not number(value): + return '—' + for unit in ('B', 'KiB', 'MiB', 'GiB', 'TiB'): + if value < 1024 or unit == 'TiB': + return decimal(value, f' {unit}') + value /= 1024 + + +def capacities(path): + # Quotas are metadata, not recomputed CPU statistics. Bound reads on long runs. + result = {} + try: + with path.open() as source: + for index, line in enumerate(source): + if index == 256: + break + record = json.loads(line) + cores = record.get('host', {}).get('cpu_count') + if number(cores) and cores > 0: + result.setdefault('host', cores) + for role, values in record.get('containers', {}).items(): + quota = values.get('configured_cpu_limit_cores') + if number(quota): + result.setdefault(role, quota) + if all(role in result for role in ('host', 'runtime', 'db')): + break + except (OSError, ValueError, TypeError, AttributeError): + pass + return result + + +def render_resources(scopes, table): + rows = [] + for scope, directory in scopes: + try: + summary = json.loads((directory / 'resources.summary.json').read_text()) + phases = summary.get('phases', {}) if summary.get('schema') == 'o11ybench-resources-v1' else {} + except (OSError, ValueError, TypeError, AttributeError): + phases = {} + if not isinstance(phases, dict) or not phases: + rows.append([scope, 'unavailable', *['—'] * 9]) + continue + limits = capacities(directory / 'resources.jsonl') + for phase, data in phases.items(): + for role in ('host', 'runtime', 'db'): + values = data.get('host', {}) if role == 'host' else data.get('containers', {}).get(role, {}) + if role == 'db' and scope == 'dataset': + continue + quota = limits.get(role) + cpu_samples = values.get('cpu_used_cores', {}).get('samples') + rows.append([scope, phase, role, cpu_samples, + decimal(metric(values, 'cpu_used_cores', 'mean')), + decimal(metric(values, 'cpu_used_cores', 'max')), + 'unlimited' if quota == 0 else decimal(quota), + memory(metric(values, 'memory_used_bytes' if role == 'host' else 'memory_usage_bytes', 'max')), + decimal(metric(values, 'io_wait_percent', 'mean'), '%'), + decimal(metric(values, 'io_wait_percent', 'max'), '%'), + memory(metric(values, 'swap_used_bytes', 'max'))]) + return '\n'.join(['## Sampled resources', '', + 'CPU is measured in cores. Mean/max summarize available samples, not continuous peaks; ' + 'short phases may have no CPU sample. CPU capacity is host cores or the configured container quota. ' + 'Host memory is total minus available; Docker memory excludes cache where supported. ' + 'Swap is host-wide usage, not swap activity. Missing evidence stays —.', '', + table(['Scope', 'Phase', 'Component', 'CPU samples', 'Mean CPU cores', 'Max CPU cores', + 'CPU capacity', 'Max memory', 'Mean I/O wait', 'Max I/O wait', 'Max swap'], rows), '']) diff --git a/.github/scripts/check-nightly-release.sh b/.github/scripts/check-nightly-release.sh new file mode 100755 index 00000000000..82eb910d6fd --- /dev/null +++ b/.github/scripts/check-nightly-release.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash + +set -euo pipefail + +# Returns whether the current release run should proceed, 'true' or 'false'. +# +# The nightly schedule runs every weekday, but a scheduled release only proceeds +# when the latest published nightly release is older than +# NIGHTLY_RELEASE_MAX_AGE_DAYS days. This way, a failed nightly run is retried +# by the next weekday schedule, instead of leaving a whole week without nightly +# builds. Non-scheduled runs (tag pushes and manual dispatches) always proceed. +# +# Only the result is written to stdout, so that it can be used in GitHub Actions +# outputs; all diagnostics go to stderr. +# +# You can run as following examples: +# GITHUB_EVENT_NAME=push ./check-nightly-release.sh +# GITHUB_EVENT_NAME=schedule GITHUB_REPOSITORY=GreptimeTeam/greptimedb GH_TOKEN=ghp_xxx ./check-nightly-release.sh +# GITHUB_EVENT_NAME=schedule NIGHTLY_RELEASE_MAX_AGE_DAYS=5 GITHUB_REPOSITORY=GreptimeTeam/greptimedb GH_TOKEN=ghp_xxx ./check-nightly-release.sh +if [ -z "$GITHUB_EVENT_NAME" ]; then + echo "GITHUB_EVENT_NAME is empty" >&2 + exit 1 +fi + +if [ "$GITHUB_EVENT_NAME" != "schedule" ]; then + echo "Not a scheduled run, the nightly check is skipped." >&2 + echo "true" + exit 0 +fi + +MAX_AGE_DAYS="${NIGHTLY_RELEASE_MAX_AGE_DAYS:-5}" +if ! [[ "$MAX_AGE_DAYS" =~ ^[0-9]+$ ]]; then + echo "Invalid NIGHTLY_RELEASE_MAX_AGE_DAYS: ${MAX_AGE_DAYS}" >&2 + exit 1 +fi + +SECONDS_PER_DAY=86400 + +# The publish time of the latest published nightly release, like the tag +# 'v0.2.0-nightly-20230313'. Fail open on GitHub API errors: proceeding with +# the release is safer than skipping it. +if ! latest_nightly_published_at="$(gh api "repos/${GITHUB_REPOSITORY}/releases?per_page=100" \ + --jq '[.[] | select(.draft == false and (.tag_name | contains("nightly")))][0].published_at // empty')"; then + echo "::warning::Failed to query the latest nightly release from the GitHub API, will proceed with the nightly release." >&2 + echo "true" + exit 0 +fi + +if [ -z "$latest_nightly_published_at" ]; then + echo "No published nightly release found, will proceed with the nightly release." >&2 + echo "true" + exit 0 +fi + +published_at_epoch="$(date --date="$latest_nightly_published_at" +%s)" +now_epoch="$(date +%s)" +age_seconds=$(( now_epoch - published_at_epoch )) +age_days=$(( age_seconds / SECONDS_PER_DAY )) + +echo "The latest nightly release was published at ${latest_nightly_published_at} (${age_days} days ago)." >&2 + +if [ "$age_seconds" -gt $(( MAX_AGE_DAYS * SECONDS_PER_DAY )) ]; then + echo "The latest nightly release is older than ${MAX_AGE_DAYS} days, will proceed with the nightly release." >&2 + echo "true" +else + echo "The latest nightly release is not older than ${MAX_AGE_DAYS} days, will skip the nightly release." >&2 + echo "false" +fi diff --git a/.github/scripts/ci-runner-config.py b/.github/scripts/ci-runner-config.py new file mode 100644 index 00000000000..a242b263c3e --- /dev/null +++ b/.github/scripts/ci-runner-config.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resolve dispatch provider and instance type before allocating a runner.""" + +import argparse +import re + +from runner_utils import append_github_output + +DEFAULT_INSTANCE_TYPES = {"Aliyun": "ecs.c9i.2xlarge", "AWS": "c7i.2xlarge"} + + +def resolve_instance_type(provider: str, requested: str) -> str: + if provider not in DEFAULT_INSTANCE_TYPES: + raise ValueError("provider must be Aliyun or AWS") + instance_type = DEFAULT_INSTANCE_TYPES[provider] if requested == "auto" else requested + pattern = r"ecs\.[a-z0-9-]+\.[a-z0-9-]+" if provider == "Aliyun" else r"[a-z0-9-]+\.[a-z0-9-]+" + if not re.fullmatch(pattern, instance_type): + raise ValueError(f"Invalid {provider} instance type: {instance_type!r}") + return instance_type + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--provider", required=True) + parser.add_argument("--instance-type", default="auto") + args = parser.parse_args() + instance_type = resolve_instance_type(args.provider, args.instance_type) + append_github_output("instance_type", instance_type) + print(f"CI runner: {args.provider} / {instance_type}") + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/ci-runners.md b/.github/scripts/ci-runners.md new file mode 100644 index 00000000000..6e4b2b3a4f7 --- /dev/null +++ b/.github/scripts/ci-runners.md @@ -0,0 +1,73 @@ +# CI runners + +The agent observability and long-range benchmark workflows select `Aliyun` or +`AWS` through the `provider` dispatch input. `instance_type: auto` selects +`ecs.c9i.2xlarge` or `c7i.2xlarge`, respectively (8 vCPU, 16 GiB). +Existing build/release and JSONBench workflows keep their current runner actions. + +Aliyun uses the existing ECS actions and image. Its provision and teardown +scripts share only GitHub API and output helpers from `runner_utils.py` with +the new AWS scripts. + +AWS uses `aws-ec2-runner-provision.py` and `aws-ec2-runner-teardown.py` directly. +Repository variables supply `EC2_RUNNER_REGION`, `EC2_RUNNER_SUBNET_ID`, and +`EC2_RUNNER_SECURITY_GROUP_ID`; credentials use `AWS_ACCESS_KEY_ID`, +`AWS_SECRET_ACCESS_KEY`, and `GH_PERSONAL_ACCESS_TOKEN` secrets. +The current region is Singapore (`ap-southeast-1`). By default the script selects +the newest available Canonical Ubuntu 24.04 amd64 server AMI in that region +(owner `099720109477`) and logs its ID. Set `BENCHMARK_EC2_IMAGE_ID` to pin a +prepared Ubuntu AMI. Existing build CI's `EC2_RUNNER_LINUX_AMD64_IMAGE_ID` is +unaffected. Both workflows default `system_disk_gib` to `auto` (also accepting +`0`): agent observability uses S=80, P=100, M=500 GiB, and long-range uses 80 GiB. +These defaults are the same for both providers. Explicit values from 20 to 2048 +GiB override them and must be at least the selected AMI's root snapshot size. +The official 24.04 image does not have the old custom AMI's 500 GiB minimum. + +AWS user data registers a fresh ephemeral runner. Git, curl, tar, and sudo are +expected in the Ubuntu server image. Docker group membership is set before the +runner service starts, and this dedicated VM grants the runner passwordless sudo +so its job can install dependencies. Like ECS, the host creates 16 GiB swap with +swappiness 10, disables systemd-oomd and background automatic updates, and sets +`OOMPolicy=continue` on the runner service. Explicit setup package installation +remains enabled. + +Both benchmark jobs call `.github/actions/setup-benchmark` after checkout. It +reuses existing tools, installs missing curl/jq/python3/rsync, installs Docker CE +if absent, and checks Docker access. Missing packages require passwordless sudo; +Aliyun's prepared image normally provides the tools and Docker already. Existing +Docker installations are preserved, so exact versions need not match between +providers. Provisioning must establish Docker permissions before starting the +runner; adding a group inside the running job would not update its groups. + +The subnet must provide outbound access to Ubuntu and Docker packages, GitHub, +and the configured container registries. The IAM identity needs EC2 +image/instance-type discovery, instance and volume tagging, launch, describe, +console output, and termination permissions. GitHub credentials need repository +runner administration access. + +Creation publishes instance identity before waiting for registration. Teardown +can recover missing outputs using repository and workflow-attempt tags. The +daily AWS janitor (03:23 UTC) removes expired instances tagged `managed-by=query-regression-ci` +and unregisters their runners. Both providers use the existing +`managed-by=query-regression-ci`, `query-regression-run-id`, and +`runner-ttl-hours` tags. AWS retains a run-attempt suffix in its run identity to +isolate retries, plus repository and runner-name tags to restrict cleanup. +Legacy build instances do not match these ownership filters. +Both janitors use a 4-hour fallback when the TTL tag is absent and skip invalid +TTL tags. Benchmark jobs explicitly tag their configured TTL (default 8 hours). +Cleanup does not check whether the GitHub job is still running. Expired instances +are removed on the next daily sweep; normal workflow teardown runs immediately. +TTL is measured from instance launch, so it must cover provisioning and the +benchmark. The scheduled janitor becomes active after landing on the default +branch. + +Local validation: + +```sh +python3 tests/perf/test_ci_runner_scripts.py +python3 tests/perf/test_aliyun_ecs_runner_scripts.py +actionlint .github/workflows/agent-observability.yml .github/workflows/vmbench-long-range.yml .github/workflows/aws-ec2-runner-janitor.yml +``` + +A real AWS dispatch is still needed to verify AMI boot, package +installation, registry connectivity, and end-to-end benchmark execution. diff --git a/.github/scripts/ci-slash.py b/.github/scripts/ci-slash.py index 61efa97fdf4..528df3feee7 100644 --- a/.github/scripts/ci-slash.py +++ b/.github/scripts/ci-slash.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Admit `/ci` comments and dispatch a pinned draft-PR CI workflow.""" +"""Admit `/ci` comments and dispatch a draft-PR CI workflow.""" import json, os, re, sys, urllib.request COMMAND = re.compile(r'^/ci(?:\s+(.+))?$') @@ -29,7 +29,7 @@ OPTIONS = { 'fuzz chaos': ('integration.yml', 'chaos', 'chaos fuzz'), 'fuzz all': ('integration.yml', 'all', 'all fuzz'), } -HELP = '''Available draft-PR CI commands:\n\n- `/ci` or `/ci rust` — Rust CI\n- `/ci integration` — integration CI without fuzz\n- `/ci checks` or `/ci docs`\n- `/ci fuzz standalone|distributed|chaos`\n- `/ci fuzz all` — all fuzz suites (admin only)\n\nCommands require repository admin permission and a same-repository open draft PR. CI is pinned to the current head SHA; comment again after a push.''' +HELP = '''Available draft-PR CI commands:\n\n- `/ci` — standard CI\n- `/ci rust` — Rust CI\n- `/ci integration` — integration CI without fuzz\n- `/ci checks` or `/ci docs`\n- `/ci fuzz standalone|distributed|chaos`\n- `/ci fuzz all` — all fuzz suites (admin only)\n\nCommands require the PR author or repository write/maintain/admin permission and an open draft PR. Fork PRs require repository write permission and rerun existing pull-request CI; fuzz commands are same-repository only. `/ci fuzz all` requires admin permission. Same-repository CI runs the branch head at dispatch time. Fork CI reruns the original pull-request revision (within GitHub rerun limits); sync your branch with the updated CI configuration first.''' def api(path): req=urllib.request.Request(os.environ['GITHUB_API_URL']+path, headers={'Authorization':'Bearer '+os.environ['GITHUB_TOKEN'],'Accept':'application/vnd.github+json'}) @@ -46,7 +46,7 @@ def reject(number, text): out(skip='true',pr_number=number,reply='CI command ign def main(): if os.environ.get('DISPATCH_SENDER') != 'github-actions[bot]': return reject('', 'invalid dispatch sender.') comment=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/issues/comments/'+os.environ['COMMENT_ID']) - body=comment.get('body','').splitlines()[0].strip(); m=COMMAND.fullmatch(body) + body=(comment.get('body') or '').partition('\n')[0].strip(); m=COMMAND.fullmatch(body) number=str(comment.get('issue_url','').rstrip('/').split('/')[-1]) if not m: return reject(number,'comment is not a `/ci` command.') arg=(m.group(1) or '').strip().lower() @@ -54,13 +54,38 @@ def main(): if arg=='help': out(skip='true',pr_number=number,reply=HELP); return 0 pr=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/pulls/'+number) if pr.get('state')!='open' or not pr.get('draft'): return reject(number,'PR must be open and draft.') - if pr.get('head',{}).get('repo',{}).get('full_name') != os.environ['GITHUB_REPOSITORY']: return reject(number,'fork PRs are not admitted.') - head=pr.get('head',{}).get('sha','') - if head != os.environ.get('DISPATCH_HEAD_SHA') or not re.fullmatch('[0-9a-f]{40}',head): return reject(number,'PR head changed; comment again.') + fork=pr.get('head',{}).get('repo',{}).get('full_name') != os.environ['GITHUB_REPOSITORY'] + if fork and arg.startswith('fuzz '): return reject(number,'fuzz commands require a same-repository PR.') + head=pr.get('head',{}) + head_sha=head.get('sha','') + head_ref=head.get('ref','') + if head_sha != os.environ.get('DISPATCH_HEAD_SHA') or not re.fullmatch('[0-9a-f]{40}',head_sha) or not head_ref: return reject(number,'PR head changed; comment again.') actor=comment.get('user',{}).get('login','') - permission=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/collaborators/'+actor+'/permission').get('user',{}).get('permission') - if permission!='admin': return reject(number,'repository admin permission is required.') + if not actor: return reject(number,'comment author is missing.') + is_author=actor==pr.get('user',{}).get('login') + if fork or arg=='fuzz all' or not is_author: + permission=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/collaborators/'+actor+'/permission').get('permission') + if arg=='fuzz all': + if permission!='admin': return reject(number,'repository admin permission is required for `/ci fuzz all`.') + elif permission not in ('write','maintain','admin'): + return reject(number,'PR author or repository write permission is required.') workflow, profile, label=OPTIONS[arg] - out(skip='false',pr_number=number,head_sha=head,workflow=workflow,fuzz_profile=profile,reply=f'Dispatched {label} for `{head}`.') + if fork: + runs=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/actions/runs?event=pull_request&head_sha='+head_sha+'&per_page=100')['workflow_runs'] + selected={} + for run in runs: + name=run.get('path','').split('@',1)[0].removeprefix('.github/workflows/') + if name not in workflow.split(',') or name in selected: continue + if run.get('head_sha') != head_sha or run.get('head_repository',{}).get('full_name') != head['repo']['full_name'] or run.get('head_branch') != head_ref: continue + selected[name]=run + if not selected: return reject(number,'no pull-request CI runs found for this head; push a new commit first.') + if arg and workflow not in selected: return reject(number,'selected CI has no pull-request run for this head.') + if any(run['status'] != 'completed' for run in selected.values()): return reject(number,'pull-request CI is still pending; approve or wait for it, then comment again.') + for run in selected.values(): + original=run if run.get('run_attempt',1)==1 else api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/actions/runs/'+str(run['id'])+'/attempts/1') + if original.get('conclusion') != 'skipped': return reject(number,'only originally skipped draft CI runs can be requested; push a new commit while draft.') + out(skip='false',pr_number=number,run_ids=','.join(str(run['id']) for run in selected.values()),reply=f'Requested rerun of {", ".join(selected)} for `{head_sha}`.') + return 0 + out(skip='false',pr_number=number,head_sha=head_sha,head_ref=head_ref,workflow=workflow,fuzz_profile=profile,reply=f'Dispatched {label} for branch `{head_ref}`.') return 0 if __name__ == '__main__': sys.exit(main()) diff --git a/.github/scripts/pull-benchmark-image.sh b/.github/scripts/pull-benchmark-image.sh new file mode 100644 index 00000000000..1a26a4aeaee --- /dev/null +++ b/.github/scripts/pull-benchmark-image.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Pull only: never retry database writes, benchmarks, or cloud provisioning. +set -euo pipefail + +if [[ $# != 1 || -z "$1" || "$1" == -* ]]; then + echo "Usage: $0 IMAGE" >&2 + exit 2 +fi +image=$1 +attempts=3 +for ((attempt=1; attempt<=attempts; attempt++)); do + echo "Pulling $image (attempt $attempt/$attempts)" >&2 + if timeout --kill-after=10s 180s docker pull "$image"; then + exit 0 + else + status=$? + fi + if ((attempt == attempts)); then + echo "Failed to pull $image after $attempts attempts (exit $status)" >&2 + exit "$status" + fi + delay=$((attempt * 10)) + echo "Pull failed (exit $status); retrying in ${delay}s" >&2 + sleep "$delay" +done diff --git a/.github/scripts/run-fuzz-targets-test.sh b/.github/scripts/run-fuzz-targets-test.sh index 5e4cda64366..9705b6bc5af 100755 --- a/.github/scripts/run-fuzz-targets-test.sh +++ b/.github/scripts/run-fuzz-targets-test.sh @@ -31,7 +31,7 @@ new_fixture() { #!/usr/bin/env bash set -euo pipefail printf '%s\t%s\t%s\n' "${GT_FUZZ_DUMP_DIR}" "$*" "${MOCK_CARGO_MARKER:-}" >>"${MOCK_CARGO_LOG}" -target="$3" +target="$4" case " ${MOCK_FAIL_TARGETS:-} " in *" ${target} "*) exit 17 ;; esac @@ -117,6 +117,20 @@ EOF set -e } +test_rustup_toolchain_env_is_honored() { + new_fixture + export RUSTUP_TOOLCHAIN=nightly-2026-03-21 + run_fixture $'fuzz_create_table' true true "fuzz_create_table" + unset RUSTUP_TOOLCHAIN + + assert_eq 1 "${fixture_status}" "pinned toolchain run status" + grep -q -- '+nightly-2026-03-21 fuzz run fuzz_create_table' "${fixture}/cargo.log" || \ + fail "pinned toolchain missing from cargo invocation" + grep -q 'cargo +nightly-2026-03-21 fuzz run fuzz_create_table' "${fixture}/artifacts/summary.md" || \ + fail "pinned toolchain missing from reproduce command" + cleanup_fixture +} + test_successful_targets_run_in_order() { new_fixture run_fixture $'fuzz_create_table\nfuzz_insert' false true "" @@ -125,7 +139,7 @@ test_successful_targets_run_in_order() { assert_eq 2 "$(wc -l <"${fixture}/cargo.log" | tr -d ' ')" "cargo invocation count" assert_eq \ $'fuzz_create_table\nfuzz_insert' \ - "$(awk -F '\t' '{print $2}' "${fixture}/cargo.log" | sed -E 's/^fuzz run ([^ ]+).*/\1/')" \ + "$(awk -F '\t' '{print $2}' "${fixture}/cargo.log" | sed -E 's/^\+nightly fuzz run ([^ ]+).*/\1/')" \ "target order" grep -q -- '--features=unstable' "${fixture}/cargo.log" || fail "unstable feature missing" grep -q -- '-max_total_time=120' "${fixture}/cargo.log" || fail "fuzz time missing" @@ -175,7 +189,7 @@ test_fail_fast_stops_after_first_failure() { fail "skipped target annotation missing" grep -q '### Reproduce failed targets' "${fixture}/artifacts/summary.md" || \ fail "reproduction section missing" - grep -q 'cargo fuzz run fuzz_insert' "${fixture}/artifacts/summary.md" || \ + grep -q 'cargo +nightly fuzz run fuzz_insert' "${fixture}/artifacts/summary.md" || \ fail "reproduction command missing" cleanup_fixture } @@ -367,5 +381,6 @@ test_collector_keeps_target_scopes_separate test_cluster_collector_honors_target_scope_and_namespace test_setup_failure_writes_artifact_contract test_setup_failure_keeps_manifest_when_collection_fails +test_rustup_toolchain_env_is_honored printf 'All fuzz orchestration script tests passed.\n' diff --git a/.github/scripts/run-fuzz-targets.sh b/.github/scripts/run-fuzz-targets.sh index 226a393e16c..205ba58c419 100755 --- a/.github/scripts/run-fuzz-targets.sh +++ b/.github/scripts/run-fuzz-targets.sh @@ -104,7 +104,7 @@ write_summary() { reproduce_args+=("GT_FUZZ_INSTANCE_ROOT_DIR=${GT_FUZZ_INSTANCE_ROOT_DIR}") fi fi - reproduce_args+=(cargo fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) + reproduce_args+=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) if [[ "${fuzz_unstable}" == true ]]; then reproduce_args+=(--features=unstable) fi @@ -195,7 +195,7 @@ for ((index = 0; index < ${#targets[@]}; index++)); do fi run_args=("${fuzz_binary}" "-max_total_time=${FUZZ_MAX_TOTAL_TIME}" "-artifact_prefix=${target_dir}/libfuzzer/") else - run_args=(cargo fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) + run_args=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) if [[ "${fuzz_unstable}" == true ]]; then run_args+=(--features=unstable) fi diff --git a/.github/scripts/runner_utils.py b/.github/scripts/runner_utils.py new file mode 100644 index 00000000000..afd827f62cf --- /dev/null +++ b/.github/scripts/runner_utils.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small shared GitHub runner API and workflow-output helpers.""" + +from __future__ import annotations + +import json +import os +import urllib.error +import urllib.request + + +def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") if body is not None else None + request = urllib.request.Request( + f"https://api.github.com{path}", + data=data, + method=method, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + if response.status == 204: + return {} + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + # GitHub's error body says exactly why (e.g. "Must have admin rights to + # Repository" for a PAT without the required scope); surface it instead + # of a bare "HTTP Error 403". + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" + "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " + "'repo' scope (classic PAT) or 'Administration: write' on the " + "repository (fine-grained PAT)." + ) from error + + +def create_registration_token(github_token: str, repo: str) -> str: + response = github_api( + github_token, + "POST", + f"/repos/{repo}/actions/runners/registration-token", + body={}, + ) + return response["token"] + + +def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: + page = 1 + while True: + response = github_api( + github_token, + "GET", + f"/repos/{repo}/actions/runners?per_page=100&page={page}", + ) + runners = response.get("runners", []) + for runner in runners: + if runner.get("name") == name: + return runner + if len(runners) < 100: + return None + page += 1 + + +def append_github_output(name: str, value: str) -> None: + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(f"{name}={value}\n") + + +def deregister_runner(token: str, repo: str, runner_name: str) -> bool: + runner = find_runner_by_name(token, repo, runner_name) + if runner is None: + print(f"Runner {runner_name} is not registered", flush=True) + return True + try: + github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") + print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) + return True + except Exception as error: # noqa: BLE001 + print(f"Failed to deregister runner {runner_name}: {error}", flush=True) + return False diff --git a/.github/scripts/tests/test_ci_slash.py b/.github/scripts/tests/test_ci_slash.py index afcf8031cdb..c8c01ceda15 100644 --- a/.github/scripts/tests/test_ci_slash.py +++ b/.github/scripts/tests/test_ci_slash.py @@ -14,8 +14,11 @@ # limitations under the License. import importlib.util +import json import os import tempfile +import subprocess +import textwrap import unittest from pathlib import Path from unittest.mock import patch @@ -27,6 +30,9 @@ ci = importlib.util.module_from_spec(spec) spec.loader.exec_module(ci) +OPTIONS_TO_TEST = ["/ci", "/ci rust", "/ci fuzz chaos", "/ci fuzz all"] + + class CiSlashTest(unittest.TestCase): def setUp(self): self.output = tempfile.NamedTemporaryFile(delete=False) @@ -54,23 +60,145 @@ class CiSlashTest(unittest.TestCase): self.assertIn("skip=true", output) self.assertIn("Available draft-PR CI commands", output) - def test_dispatches_full_suite_at_verified_head(self): + def test_dispatches_full_suite_with_top_level_admin_permission(self): output = self.run_main([ {"body": "/ci", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "admin"}}, - {"state": "open", "draft": True, "head": {"sha": "a" * 40, "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, - {"user": {"permission": "admin"}}, + {"state": "open", "draft": True, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + {"permission": "admin", "user": {"login": "admin"}}, ]) self.assertIn("skip=false", output) + self.assertIn("head_ref=fix/draft-ci", output) self.assertIn("workflow=rust.yml,integration.yml,checks.yml,docs.yml", output) def test_rejects_non_draft_before_permission_lookup(self): output = self.run_main([ {"body": "/ci fuzz chaos", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "admin"}}, - {"state": "open", "draft": False, "head": {"sha": "a" * 40, "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + {"state": "open", "draft": False, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, ]) self.assertIn("skip=true", output) self.assertIn("PR must be open and draft", output) + def test_rejects_missing_head_ref_before_permission_lookup(self): + output = self.run_main([ + {"body": "/ci", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "admin"}}, + {"state": "open", "draft": True, "head": {"sha": "a" * 40, "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + ]) + self.assertIn("skip=true", output) + self.assertIn("PR head changed; comment again.", output) + + def test_permission_matrix(self): + for command in OPTIONS_TO_TEST: + for author in (True, False): + for permission in ("admin", "maintain", "write", "triage", "read", "none", None): + with self.subTest(command=command, author=author, permission=permission): + Path(self.output.name).write_text("") + responses = [ + {"body": command, "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "actor"}}, + {"state": "open", "draft": True, "user": {"login": "actor" if author else "other"}, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + ] + if command == "/ci fuzz all" or not author: + responses.append({"permission": permission}) + output = self.run_main(responses) + allowed = permission == "admin" if command == "/ci fuzz all" else author or permission in ("admin", "maintain", "write") + self.assertIn("skip=false" if allowed else "skip=true", output) + + def test_reply_uses_issue_comment_endpoint(self): + workflow = SCRIPT.parents[1] / "workflows" / "ci-slash.yml" + step = workflow.read_text().split(" - name: Reply with command result\n", 1)[1] + command = textwrap.dedent(step.split(" run: |\n", 1)[1]) + reply = "Denied: `example`\nSecond line with $variables and 'quotes'" + with tempfile.TemporaryDirectory() as directory: + gh = Path(directory) / "gh" + gh.write_text("#!/usr/bin/env python3\nimport json, sys\nprint(json.dumps(sys.argv[1:]))\n") + gh.chmod(0o755) + env = {**os.environ, "PATH": directory + os.pathsep + os.environ["PATH"], "GITHUB_REPOSITORY": "GreptimeTeam/greptimedb", "PR_NUMBER": "42", "REPLY": reply, "ADMIT_OUTCOME": "success", "DISPATCH_OUTCOME": "success"} + result = subprocess.run(["bash", "-eu", "-c", command], env=env, check=True, capture_output=True, text=True) + self.assertEqual(json.loads(result.stdout), ["api", "--method", "POST", "/repos/GreptimeTeam/greptimedb/issues/42/comments", "-f", "body=" + reply]) + + def test_failure_reply_and_partial_dispatch(self): + workflow = (SCRIPT.parents[1] / "workflows" / "ci-slash.yml").read_text() + dispatch = textwrap.dedent(workflow.split(" - name: Dispatch selected CI workflow\n", 1)[1].split(" run: |\n", 1)[1].split(" - name: Reply", 1)[0]) + reply = textwrap.dedent(workflow.split(" - name: Reply with command result\n", 1)[1].split(" run: |\n", 1)[1]) + with tempfile.TemporaryDirectory() as directory: + gh = Path(directory) / "gh" + log = Path(directory) / "calls" + gh.write_text("#!/usr/bin/env python3\nimport json, os, sys\nwith open(os.environ['CALLS'], 'a') as f: f.write(json.dumps(sys.argv[1:]) + '\\n')\nsys.exit(1 if '/runs/8/rerun' in sys.argv[-1] else 0)\n") + gh.chmod(0o755) + env = {**os.environ, "PATH": directory + os.pathsep + os.environ["PATH"], "CALLS": str(log), "GITHUB_REPOSITORY": "GreptimeTeam/greptimedb", "RUN_IDS": "7,8,9", "PR_NUMBER": "42", "REPLY": "Success", "ADMIT_OUTCOME": "success", "DISPATCH_OUTCOME": "failure", "GITHUB_SERVER_URL": "https://github.com", "GITHUB_RUN_ID": "123"} + result = subprocess.run(["bash", "-eu", "-c", dispatch], env=env, capture_output=True) + self.assertNotEqual(result.returncode, 0) + subprocess.run(["bash", "-eu", "-c", reply], env=env, check=True, capture_output=True) + calls = [json.loads(line) for line in log.read_text().splitlines()] + self.assertEqual([call[-1] for call in calls[:2]], ["/repos/GreptimeTeam/greptimedb/actions/runs/7/rerun", "/repos/GreptimeTeam/greptimedb/actions/runs/8/rerun"]) + self.assertIn("CI trigger failed; some workflows may already have been requested", calls[2][-1]) + self.assertIn("/actions/runs/123", calls[2][-1]) + self.assertIn("!cancelled()", workflow.split(" - name: Reply with command result", 1)[1]) + + def test_fork_rerun_gates_cover_standard_jobs(self): + import re + for name, count in [("rust", 7), ("integration", 5), ("checks", 5), ("docs", 3)]: + workflow = (SCRIPT.parents[1] / "workflows" / (name + ".yml")).read_text() + gates = re.findall(r"^ if:.*github.run_attempt > 1.*$", workflow, re.MULTILINE) + self.assertEqual(len(gates), count, name) + self.assertTrue(all("always()" not in gate for gate in gates)) + + def test_command_permissions(self): + workflow = SCRIPT.parents[1] / "workflows" / "slash-command-dispatch.yml" + config = workflow.read_text().split("config: >-", 1)[1].split(" - name:", 1)[0] + self.assertEqual(json.loads(config), [ + {"command": "query-regression", "permission": "admin", "issue_type": "pull-request"}, + {"command": "ci", "permission": "none", "issue_type": "pull-request"}, + ]) + + def test_fork_requires_writer_even_for_author(self): + for permission in ("read", "write", "maintain", "admin"): + with self.subTest(permission=permission): + Path(self.output.name).write_text("") + responses = [ + {"body": "/ci", "issue_url": "https://api.github.test/issues/42", "user": {"login": "actor"}}, + {"state": "open", "draft": True, "user": {"login": "actor"}, "head": {"sha": "a" * 40, "ref": "fork-branch", "repo": {"full_name": "actor/greptimedb"}}}, + {"permission": permission}, + ] + if permission != "read": + responses.append({"workflow_runs": [ + {"id": 7, "path": ".github/workflows/rust.yml", "head_sha": "a" * 40, "head_branch": "fork-branch", "head_repository": {"full_name": "actor/greptimedb"}, "status": "completed", "conclusion": "skipped"}, + {"id": 8, "path": ".github/workflows/checks.yml", "head_sha": "a" * 40, "head_branch": "fork-branch", "head_repository": {"full_name": "someone/greptimedb"}, "status": "completed", "conclusion": "skipped"}, + ]}) + output = self.run_main(responses) + if permission == "read": + self.assertIn("skip=true", output) + else: + self.assertIn("run_ids=7\n", output) + self.assertNotIn("head_ref=", output) + + def test_fork_rejects_non_draft_original_run(self): + output = self.run_main([ + {"body": "/ci rust", "issue_url": "https://api.github.test/issues/42", "user": {"login": "writer"}}, + {"state": "open", "draft": True, "head": {"sha": "a" * 40, "ref": "fork-branch", "repo": {"full_name": "actor/greptimedb"}}}, + {"permission": "write"}, + {"workflow_runs": [{"id": 7, "path": ".github/workflows/rust.yml@main", "head_sha": "a" * 40, "head_branch": "fork-branch", "head_repository": {"full_name": "actor/greptimedb"}, "status": "completed", "conclusion": "success", "run_attempt": 2}]}, + {"conclusion": "success"}, + ]) + self.assertIn("only originally skipped draft CI runs", output) + self.assertIn("skip=true", output) + + def test_author_cannot_bypass_pr_guards(self): + for change, reason in [ + ({"state": "closed"}, "PR must be open and draft"), + ({"draft": False}, "PR must be open and draft"), + ({"head": {"sha": "b" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, "PR head changed"), + ]: + with self.subTest(change=change): + Path(self.output.name).write_text("") + pr = {"state": "open", "draft": True, "user": {"login": "actor"}, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}} + pr.update(change) + output = self.run_main([ + {"body": "/ci rust", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "actor"}}, + pr, + ]) + self.assertIn("skip=true", output) + self.assertIn(reason, output) + if __name__ == "__main__": unittest.main() diff --git a/.github/scripts/tracesbench-summary.py b/.github/scripts/tracesbench-summary.py new file mode 100644 index 00000000000..087a67498f4 --- /dev/null +++ b/.github/scripts/tracesbench-summary.py @@ -0,0 +1,169 @@ +# 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. + +import argparse +import json +import html +import math +from pathlib import Path + +from benchmark_resource_summary import render_resources + +TARGETS = ('greptimedb', 'victoriatraces', 'tempo') + + +def duration(ms): + if not isinstance(ms, (int, float)) or not math.isfinite(ms) or ms < 0: + return '—' + if ms < 1: + return f'{ms * 1000:.0f} µs' + if ms < 1000: + return f'{ms:.1f} ms' + if ms < 60_000: + return f'{ms / 1000:.2f} s' + minutes, seconds = divmod(round(ms / 1000), 60) + hours, minutes = divmod(minutes, 60) + return f'{hours}h {minutes}m {seconds}s' if hours else f'{minutes}m {seconds}s' + + +def cell(value): + return html.escape(str(value if value is not None else '—')).replace('|', '|').replace('\n', ' ') + + +def read(path): + if not path.exists(): + return {} + return json.loads(path.read_text()) + + +def table(headers, rows): + return '\n'.join(['| ' + ' | '.join(map(cell, headers)) + ' |', + '| ' + ' | '.join('---' for _ in headers) + ' |'] + + ['| ' + ' | '.join(map(cell, row)) + ' |' for row in rows]) + + +def delta(value, baseline): + if not all(isinstance(v, (int, float)) and math.isfinite(v) for v in (value, baseline)): + return '—' + return f'{(value / baseline - 1) * 100:+.1f}%' if baseline > 0 and value >= 0 else '—' + + + +def load(root): + report = read(root / 'report.json') + runs = [read(p) for p in sorted((root / 'runs').glob('*/run.json'))] + return report, runs + + +def benchmark(report, runs, targets=TARGETS): + lines = ['# Tracesbench benchmark summary', '', + 'Median Δ = (target / GreptimeDB − 1) × 100%; negative is faster. ' + 'Min/median/max come from the validated report, not recalculated samples. ' + 'Counts exclude warmup. Missing or excluded results have no latency comparison.', + '', f'Classification: {cell(report.get("classification"))}. ' + 'See the uploaded HTML for query semantics and client-side compensation.', ''] + statistics = {} + for corpus in report.get('per_corpus', []): + for stat in corpus.get('statistics', []): + if stat['target'] not in targets: + continue + statistics[(corpus['corpus'], stat['round'], stat['target'], stat['query_id'])] = stat + rows, details = [], [] + for target in targets: + selected = [run for run in runs if run.get('target') == target] + if not selected: + lines.append(f'- **{target}**: unavailable / not run.') + for run in selected: + lines.append(f'- **{target}**: lifecycle ' + f'{"completed" if run.get("success") is True else "failed / incomplete"}; ' + f'query failures: {len(run.get("query_failures", []))}.') + for qid in (f'T{i:02}' for i in range(1, 9)): + correct = [q for q in run.get('correctness', []) if q.get('query_id') == qid] + measured = [q for q in run.get('samples', []) + if q.get('query_id') == qid and q.get('phase') == 'measured'] + errors = [q for q in correct + run.get('samples', []) + if q.get('query_id') == qid and q.get('outcome') not in ('passed', 'not_executed')] + reason = (run.get('exclusions', {}).get(qid) or + run.get('execution', {}).get('query_skips', {}).get(qid)) + stat = statistics.get((run.get('corpus'), run.get('round'), target, qid), {}) + base = statistics.get((run.get('corpus'), run.get('round'), 'greptimedb', qid), {}) + status = ('failed' if errors else 'excluded' if reason else + 'passed' if stat else + 'correctness passed; report unavailable' if correct and + all(q.get('outcome') == 'passed' for q in correct) else 'unavailable / incomplete') + rows.append([run.get('corpus'), run.get('round'), qid, target, status, + duration(stat.get('min_ms')), duration(stat.get('median_ms')), + duration(stat.get('max_ms')), + delta(stat.get('median_ms'), base.get('median_ms')) if target != 'greptimedb' else '—', + sum(q.get('outcome') == 'passed' for q in measured), + sum(q.get('outcome') != 'passed' for q in measured)]) + if reason or errors: + messages = [str(q.get('error') or q.get('outcome')) for q in errors] + details.append(f'- **{target} / {qid}**: {cell(str(reason or "") + " " + "; ".join(messages))[:4000]}') + lines += ['', table(['Dataset', 'Round', 'Query', 'DB', 'Status', 'Min', 'Median', 'Max', + 'Median Δ', 'Measured success', 'Measured errors'], rows), '', + '
', 'Query failures and exclusions', '', *details, '', '
'] + if not report: + lines += ['', 'Validated report unavailable; raw outcomes are shown without latency statistics.'] + return '\n'.join(lines) + + +def lifecycle(report, runs, root, targets=TARGETS): + rows = [] + for target in targets: + selected = [r for r in runs if r.get('target') == target] or [{}] + for run in selected: + load = run.get('load', {}) + ingest = load.get('ingest', {}) + if target == 'tempo' and 'ingest' not in load: + ingest = load.get('transport', {}).get('ingest', {}) + health = run.get('post_load_health', {}) + execution = run.get('execution', {}) + start, finish = run.get('started_unix_ns'), run.get('finished_unix_ns') + elapsed = (finish - start) / 1_000_000 if isinstance(start, int) and isinstance(finish, int) else None + seconds = ingest.get('wall_seconds') + rows.append([target, run.get('corpus'), run.get('load_summary', {}).get('after'), + duration(elapsed), duration(seconds * 1000 if isinstance(seconds, (int, float)) else None), + ingest.get('successful_requests'), ingest.get('failed_requests'), + execution.get('load_workers_requested'), execution.get('max_in_flight'), + execution.get('db_cpus'), execution.get('db_memory'), + health.get('running'), health.get('oom_killed'), + run.get('cleanup', {}).get('verified_absent')]) + return '\n'.join(['# Tracesbench load / lifecycle summary', '', + 'Total is the recorded target lifecycle; ingest is the sender wall time, ' + 'not the entire load/readiness phase. Health is the post-load snapshot. ' + 'Missing evidence remains blank.', '', + table(['DB', 'Dataset', 'Loaded spans', 'Total', 'Ingest', + 'Write success', 'Write errors', 'Load workers', 'Max in flight', + 'DB CPU', 'DB memory', 'Running after load', 'OOM after load', + 'Container removed'], rows), '', + render_resources([('dataset', root / 'runs' / 'generate')] + + [(t, root / 'runs' / t) for t in targets], table)]) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--root', type=Path, required=True) + parser.add_argument('--section', choices=('benchmark', 'lifecycle'), required=True) + parser.add_argument('--targets', default=','.join(TARGETS), help='Comma-separated selected engines') + args = parser.parse_args() + targets = args.targets.split(',') + if not targets or any(t not in TARGETS for t in targets) or len(set(targets)) != len(targets): + parser.error('targets must be a non-empty comma-separated subset of greptimedb,victoriatraces,tempo') + report, runs = load(args.root) + print(benchmark(report, runs, targets) if args.section == 'benchmark' else lifecycle(report, runs, args.root, targets)) + + +if __name__ == '__main__': + main() diff --git a/.github/workflows/agent-observability.yml b/.github/workflows/agent-observability.yml index 02c4ccb81bc..d265801d4a0 100644 --- a/.github/workflows/agent-observability.yml +++ b/.github/workflows/agent-observability.yml @@ -26,8 +26,18 @@ on: - M default: S required: true + compare_greptimedb: + description: Compare baseline and candidate sequentially on the same host (adds one full GreptimeDB run) + type: boolean + default: false + required: false + baseline_greptimedb_tag: + description: Baseline Docker Hub greptime/greptimedb tag; required when comparison is enabled + type: string + default: v1.2.1 + required: false greptimedb_tag: - description: Docker Hub greptime/greptimedb tag + description: Candidate Docker Hub greptime/greptimedb tag type: string default: latest required: true @@ -46,19 +56,25 @@ on: type: string required: false runtime_image: - description: Aliyun runtime image (tag or digest reference) + description: 'Runtime image: auto = AWS ECR Public or Aliyun ACR by provider; explicit tag/digest overrides' type: string - default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260917-8111034e@sha256:ce9dee2abb1e70f907a09ea01854904a91cc97748b0c4ef6f33d9ec34044bec0 + default: auto required: true - ecs_instance_type: - description: ECS instance type (independent of Query Regression) + provider: + description: Cloud provider for the benchmark runner + type: choice + options: [Aliyun, AWS] + default: Aliyun + required: true + instance_type: + description: 'Instance type; auto = Aliyun: ecs.c9i.2xlarge, AWS: c7i.2xlarge (both 8 vCPU / 16 GiB). Enter a specific type to override.' type: string - default: ecs.c9i.2xlarge + default: auto required: true system_disk_gib: - description: System disk GiB; 0 = auto (S 80, P 100, M 500), or 20..2048 - type: number - default: 0 + description: 'System disk GiB: auto (or 0) = S 80, P 100, M 500; enter 20..2048 to override' + type: string + default: auto required: true benchmark_timeout_minutes: description: Benchmark job timeout in minutes (1..360) @@ -83,13 +99,16 @@ on: permissions: contents: read env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + BASELINE_GREPTIMEDB_TAG: ${{ inputs.baseline_greptimedb_tag }} PROFILE: ${{ inputs.dataset }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} CLICKHOUSE_TAG: ${{ inputs.clickhouse_tag }} VICTORIALOGS_TAG: ${{ inputs.victorialogs_tag }} O11YBENCH_REF: ${{ inputs.o11ybench_ref }} - RUNTIME_IMAGE: ${{ inputs.runtime_image }} - ECS_INSTANCE_TYPE: ${{ inputs.ecs_instance_type }} + RUNTIME_IMAGE: ${{ inputs.runtime_image != 'auto' && inputs.runtime_image || inputs.provider == 'AWS' && 'public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime:2f40fb3fdbdc@sha256:8940782216bf5d066308fd71f6e346d950073c41cfb5756c5f3dc8aacf069cbc' || 'greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:2f40fb3fdbdc@sha256:8940782216bf5d066308fd71f6e346d950073c41cfb5756c5f3dc8aacf069cbc' }} + PROVIDER: ${{ inputs.provider }} + REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} @@ -100,6 +119,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 10 outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} targets: ${{ steps.inputs.outputs.targets }} system_disk_gib: ${{ steps.inputs.outputs.system_disk_gib }} greptimedb: ${{ steps.inputs.outputs.greptimedb }} @@ -110,10 +130,30 @@ jobs: with: path: greptimedb persist-credentials: false - - name: Test ECS lifecycle helpers + - name: Validate version comparison + env: + RUN_GREPTIMEDB: ${{ inputs.run_greptimedb }} + run: | + python3 greptimedb/tests/perf/test_benchmark_version_comparison.py + python3 greptimedb/.github/scripts/benchmark-version-comparison.py validate + - name: Resolve CI runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" + - name: Test CI runner helpers run: | python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py + python3 greptimedb/tests/perf/test_ci_runner_scripts.py python3 greptimedb/tests/perf/test_agent_observability_summary.py + - name: Check runtime supports shared version-comparison data + if: inputs.compare_greptimedb + run: | + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" + docker run --rm --entrypoint grep "$RUNTIME_IMAGE" -q -- '--data-root)' /opt/o11ybench/scripts/run-agent-observability.sh \ + || { echo 'Comparison requires a runtime with --data-root support'; exit 1; } + - name: Validate dataset admission in runtime + run: | + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" + docker run --rm "$RUNTIME_IMAGE" scripts/logbench_agent_observability_measured.py --preflight-profile "$PROFILE" - name: Validate requested benchmark id: inputs shell: bash @@ -124,14 +164,15 @@ jobs: run: | set -euo pipefail # Validate cloud lifecycle budgets before provisioning; runtime owns workload validation. - [[ -n "$ECS_INSTANCE_TYPE" ]] case "$PROFILE" in S) auto_disk_gib=80 ;; P) auto_disk_gib=100 ;; M) auto_disk_gib=500 ;; *) echo "Unknown dataset: $PROFILE" >&2; exit 1 ;; esac - [[ "$SYSTEM_DISK_GIB" != 0 ]] || SYSTEM_DISK_GIB=$auto_disk_gib + if [[ "$SYSTEM_DISK_GIB" == auto || "$SYSTEM_DISK_GIB" == 0 ]]; then + SYSTEM_DISK_GIB=$auto_disk_gib + fi for value in "$SYSTEM_DISK_GIB" "$BENCHMARK_TIMEOUT_MINUTES" "$JANITOR_TTL_HOURS"; do [[ "$value" =~ ^[1-9][0-9]{0,3}$ ]] || { echo 'Resource budgets must be positive integers' >&2; exit 1; } done @@ -160,15 +201,16 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 45 outputs: - label: ${{ steps.ecs.outputs.label }} - instance_id: ${{ steps.ecs.outputs.instance_id }} - runner_name: ${{ steps.ecs.outputs.runner_name }} + label: ${{ steps.ecs.outputs.label || steps.aws.outputs.label }} + instance_id: ${{ steps.ecs.outputs.instance_id || steps.aws.outputs.instance_id }} + runner_name: ${{ steps.ecs.outputs.runner_name || steps.aws.outputs.runner_name }} steps: - uses: actions/checkout@v4 with: path: greptimedb persist-credentials: false - name: Create ephemeral ECS runner + if: inputs.provider == 'Aliyun' id: ecs uses: ./greptimedb/.github/actions/aliyun-ecs-create with: @@ -179,13 +221,35 @@ jobs: vswitch-id: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} security-group-id: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} - instance-type: ${{ env.ECS_INSTANCE_TYPE }} + instance-type: ${{ needs.validate.outputs.instance_type }} enable-docker: 'true' system-disk-gib: ${{ needs.validate.outputs.system_disk_gib }} ttl-hours: ${{ inputs.janitor_ttl_hours }} resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} runner-gid: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Create ephemeral EC2 runner + if: inputs.provider == 'AWS' + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} + AWS_EC2_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py benchmark: needs: - validate @@ -193,6 +257,7 @@ jobs: runs-on: ${{ needs.provision.outputs.label }} timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} env: + INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} TARGETS: ${{ needs.validate.outputs.targets }} SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} OWNER: o11ybench-${{ github.run_id }}-${{ github.run_attempt }} @@ -209,6 +274,8 @@ jobs: with: path: greptimedb persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark - name: Check Docker on ephemeral Linux host shell: bash run: | @@ -223,11 +290,11 @@ jobs: run: | set -euo pipefail actual= - [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* ]] + [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* || "$RUNTIME_IMAGE" == public.ecr.aws/* ]] images='{}' resolve_image() { local target=$1 reference=$2 variable=$3 metadata resolved revision - docker pull "$reference" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$reference" metadata=$(docker image inspect "$reference") resolved=$(jq -er '.[0].RepoDigests[0]' <<< "$metadata") if [[ "$target" == runtime ]]; then @@ -250,11 +317,14 @@ jobs: victorialogs) resolve_image "$target" "victoriametrics/victoria-logs:$VICTORIALOGS_TAG" VICTORIALOGS_IMAGE ;; esac done + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + resolve_image greptimedb-baseline "greptime/greptimedb:$BASELINE_GREPTIMEDB_TAG" BASELINE_GREPTIMEDB_IMAGE + fi jq -n --arg sha "$actual" --arg targets "$TARGETS" --arg profile "$PROFILE" \ --arg cpus "$DB_CPUS" --arg memory "$DB_MEMORY" --argjson images "$images" \ - --arg instance_type "$ECS_INSTANCE_TYPE" --argjson disk_gib "$SYSTEM_DISK_GIB" \ + --arg provider "$PROVIDER" --arg instance_type "$INSTANCE_TYPE" --argjson disk_gib "$SYSTEM_DISK_GIB" \ --argjson timeout_minutes "$BENCHMARK_TIMEOUT_MINUTES" --argjson ttl_hours "$JANITOR_TTL_HOURS" \ - '{o11ybench_sha:$sha, targets:$targets, profile:$profile, db_cpus:$cpus, db_memory:$memory, images:$images, ecs:{instance_type:$instance_type, system_disk_gib:$disk_gib, benchmark_timeout_minutes:$timeout_minutes, janitor_ttl_hours:$ttl_hours}}' > "$MANIFEST_PATH" + '{o11ybench_sha:$sha, targets:$targets, profile:$profile, db_cpus:$cpus, db_memory:$memory, images:$images, ($provider | if . == "AWS" then "ec2" else "ecs" end):{instance_type:$instance_type, system_disk_gib:$disk_gib, benchmark_timeout_minutes:$timeout_minutes, janitor_ttl_hours:$ttl_hours}}' > "$MANIFEST_PATH" - name: Checkout benchmark tools uses: actions/checkout@v4 with: @@ -273,6 +343,15 @@ jobs: --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" \ --root "$ARTIFACT_ROOT" --profile "$PROFILE" --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" cp "$MANIFEST_PATH" "$ARTIFACT_ROOT/run-manifest.json" + - name: Load and benchmark baseline GreptimeDB + id: baseline_greptimedb + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' && inputs.compare_greptimedb }} + run: |- + GITHUB_STEP_SUMMARY= bash o11ybench/scripts/run-agent-observability.sh --stage target --execute \ + --owner "$OWNER-base" --runtime-image "$RESOLVED_RUNTIME_IMAGE" --root "$ARTIFACT_ROOT/baseline" --data-root "$ARTIFACT_ROOT" \ + --profile "$PROFILE" --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + --target greptimedb --image "$BASELINE_GREPTIMEDB_IMAGE" \ + --greptimedb-http-port "$GREPTIMEDB_HTTP_PORT" --greptimedb-grpc-port "$GREPTIMEDB_GRPC_PORT" - name: Load and benchmark greptimedb id: greptimedb if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' }} @@ -318,6 +397,13 @@ jobs: GITHUB_STEP_SUMMARY= bash o11ybench/scripts/run-agent-observability.sh --stage cleanup --execute \ --owner "$OWNER" --root "$ARTIFACT_ROOT" fi + - name: Clean up baseline containers + if: always() && inputs.compare_greptimedb + run: |- + if [[ -f o11ybench/scripts/run-agent-observability.sh ]]; then + GITHUB_STEP_SUMMARY= bash o11ybench/scripts/run-agent-observability.sh --stage cleanup --execute \ + --owner "$OWNER-base" --root "$ARTIFACT_ROOT/baseline" + fi - name: Benchmark summary if: always() run: | @@ -328,6 +414,16 @@ jobs: run: | python3 greptimedb/.github/scripts/agent-observability-summary.py \ --root "$ARTIFACT_ROOT" --targets "$TARGETS" --section lifecycle >> "$GITHUB_STEP_SUMMARY" + - name: Baseline load / lifecycle summary + if: always() && inputs.compare_greptimedb + run: | + echo '## Baseline GreptimeDB' >> "$GITHUB_STEP_SUMMARY" + python3 greptimedb/.github/scripts/agent-observability-summary.py \ + --root "$ARTIFACT_ROOT/baseline" --targets greptimedb --section lifecycle >> "$GITHUB_STEP_SUMMARY" + - name: GreptimeDB version comparison + if: always() && inputs.compare_greptimedb + run: | + python3 greptimedb/.github/scripts/benchmark-version-comparison.py report --kind agent --root "$ARTIFACT_ROOT" >> "$GITHUB_STEP_SUMMARY" - name: Upload results (no dataset) if: always() uses: actions/upload-artifact@v4 @@ -343,7 +439,7 @@ jobs: needs: - provision - benchmark - if: ${{ always() && needs.provision.outputs.instance_id != '' }} + if: ${{ always() && (needs.provision.outputs.instance_id != '' || (inputs.provider == 'AWS' && needs.provision.result != 'skipped')) }} runs-on: ubuntu-latest timeout-minutes: 15 steps: @@ -352,6 +448,7 @@ jobs: path: greptimedb persist-credentials: false - name: Delete ECS and unregister runner + if: inputs.provider == 'Aliyun' uses: ./greptimedb/.github/actions/aliyun-ecs-delete with: access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} @@ -360,3 +457,20 @@ jobs: region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} instance-id: ${{ needs.provision.outputs.instance_id }} runner-name: ${{ needs.provision.outputs.runner_name }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Delete EC2 and unregister runner + if: inputs.provider == 'AWS' + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py diff --git a/.github/workflows/aws-ec2-runner-janitor.yml b/.github/workflows/aws-ec2-runner-janitor.yml new file mode 100644 index 00000000000..1c4435022b4 --- /dev/null +++ b/.github/workflows/aws-ec2-runner-janitor.yml @@ -0,0 +1,29 @@ +name: AWS EC2 Runner Janitor + +on: + schedule: + - cron: '23 3 * * *' + workflow_dispatch: + +permissions: + contents: read + +jobs: + sweep: + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + - name: Sweep expired CI instances and unregister runners + env: + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + run: uv run .github/scripts/aws-ec2-runner-teardown.py --sweep diff --git a/.github/workflows/cargo-lock-check.yml b/.github/workflows/cargo-lock-check.yml index 847be806c5f..cfea53c6750 100644 --- a/.github/workflows/cargo-lock-check.yml +++ b/.github/workflows/cargo-lock-check.yml @@ -5,6 +5,7 @@ name: Cargo.lock Diff Check # posts/updates a PR comment when the threshold is exceeded. on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] paths: - "Cargo.lock" diff --git a/.github/workflows/check-git-deps.yml b/.github/workflows/check-git-deps.yml index 6ebf79f2072..a4ec488e850 100644 --- a/.github/workflows/check-git-deps.yml +++ b/.github/workflows/check-git-deps.yml @@ -2,6 +2,7 @@ name: Check Git Dependencies on Main Branch on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] branches: [main] paths: - 'Cargo.toml' diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 6e591ba9dfd..deba5010c80 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -32,7 +32,7 @@ concurrency: jobs: check-typos-and-docs: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check typos and docs runs-on: ubuntu-latest steps: @@ -47,7 +47,7 @@ jobs: || (echo "'config/config.md' is not up-to-date, please run 'make config-docs'." && exit 1) license-header-check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} runs-on: ubuntu-latest name: Check License Header steps: @@ -63,7 +63,7 @@ jobs: config: licenserc-enterprise.toml github-script-tests: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: GitHub Script Tests runs-on: ubuntu-latest timeout-minutes: 5 @@ -97,10 +97,12 @@ jobs: python3 tests/perf/test_query_regression_nightly_refs.py python3 tests/perf/test_query_regression_slash.py python3 tests/perf/test_aliyun_ecs_runner_scripts.py + python3 tests/perf/test_ci_runner_scripts.py python3 tests/perf/test_agent_observability_summary.py + python3 tests/perf/test_tracesbench_summary.py check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check runs-on: ${{ matrix.os }} strategy: @@ -129,7 +131,7 @@ jobs: run: cargo check --locked --workspace --all-targets --all-features toml: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Toml Check runs-on: ubuntu-latest timeout-minutes: 60 diff --git a/.github/workflows/ci-slash.yml b/.github/workflows/ci-slash.yml index 1af8e60d95c..059654fd5cd 100644 --- a/.github/workflows/ci-slash.yml +++ b/.github/workflows/ci-slash.yml @@ -8,7 +8,7 @@ permissions: actions: write contents: read issues: write - pull-requests: read + pull-requests: write jobs: admit: @@ -30,23 +30,46 @@ jobs: DISPATCH_SENDER: ${{ github.event.sender.login }} DISPATCH_HEAD_SHA: ${{ github.event.client_payload.pull_request.head.sha }} run: python3 .github/scripts/ci-slash.py - - name: Reply with command result - if: ${{ steps.admit.outputs.reply != '' }} - env: - GH_TOKEN: ${{ github.token }} - REPLY: ${{ steps.admit.outputs.reply }} - PR_NUMBER: ${{ steps.admit.outputs.pr_number }} - run: printf '%s\n' "${REPLY}" | gh pr comment "${PR_NUMBER}" --body-file - - name: Dispatch selected CI workflow + id: dispatch if: ${{ steps.admit.outputs.skip == 'false' }} env: GH_TOKEN: ${{ github.token }} + RUN_IDS: ${{ steps.admit.outputs.run_ids }} WORKFLOWS: ${{ steps.admit.outputs.workflow }} - HEAD_SHA: ${{ steps.admit.outputs.head_sha }} + HEAD_REF: ${{ steps.admit.outputs.head_ref }} FUZZ_PROFILE: ${{ steps.admit.outputs.fuzz_profile }} run: | + if [[ -n "${RUN_IDS}" ]]; then + IFS=, read -ra runs <<<"${RUN_IDS}" + for run in "${runs[@]}"; do + gh api --method POST "/repos/${GITHUB_REPOSITORY}/actions/runs/${run}/rerun" + done + exit 0 + fi IFS=, read -ra workflows <<<"${WORKFLOWS}" for workflow in "${workflows[@]}"; do + inputs=(-F 'inputs[ci_command]=true') + if [[ "${workflow}" == 'integration.yml' ]]; then + inputs+=(-f "inputs[fuzz_profile]=${FUZZ_PROFILE}") + fi gh api --method POST "/repos/${GITHUB_REPOSITORY}/actions/workflows/${workflow}/dispatches" \ - -f ref="${HEAD_SHA}" -F 'inputs[ci_command]=true' -f "inputs[fuzz_profile]=${FUZZ_PROFILE}" + -f ref="${HEAD_REF}" "${inputs[@]}" done + - name: Reply with command result + if: ${{ !cancelled() && github.event.sender.login == 'github-actions[bot]' }} + env: + GH_TOKEN: ${{ github.token }} + REPLY: ${{ steps.admit.outputs.reply }} + PR_NUMBER: ${{ steps.admit.outputs.pr_number || github.event.client_payload.github.payload.issue.number }} + ADMIT_OUTCOME: ${{ steps.admit.outcome }} + DISPATCH_OUTCOME: ${{ steps.dispatch.outcome }} + # PR replies require pull-requests: write even through the issue-comment endpoint. + run: | + [[ "${PR_NUMBER}" =~ ^[0-9]+$ ]] || exit 1 + if [[ "${ADMIT_OUTCOME}" == failure || "${DISPATCH_OUTCOME}" == failure ]]; then + REPLY="CI trigger failed; some workflows may already have been requested. See ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID} for details." + fi + [[ -n "${REPLY}" ]] || exit 1 + gh api --method POST "/repos/${GITHUB_REPOSITORY}/issues/${PR_NUMBER}/comments" \ + -f "body=${REPLY}" diff --git a/.github/workflows/dependency-check.yml b/.github/workflows/dependency-check.yml index 93a366c4b08..014a41e45c6 100644 --- a/.github/workflows/dependency-check.yml +++ b/.github/workflows/dependency-check.yml @@ -2,6 +2,7 @@ name: Check Dependencies on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] branches: - main diff --git a/.github/workflows/dev-build.yml b/.github/workflows/dev-build.yml index d03fbeff14a..563435b708e 100644 --- a/.github/workflows/dev-build.yml +++ b/.github/workflows/dev-build.yml @@ -183,6 +183,10 @@ jobs: image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: ${{ inputs.large-page-size }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -217,6 +221,10 @@ jobs: image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: ${{ inputs.large-page-size }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/.github/workflows/docbot.yml b/.github/workflows/docbot.yml index 250a150c026..53033c6eb5b 100644 --- a/.github/workflows/docbot.yml +++ b/.github/workflows/docbot.yml @@ -1,10 +1,10 @@ name: Follow Up Docs on: pull_request_target: - types: [opened, edited] + types: [opened, edited, closed] concurrency: - group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + group: ${{ github.workflow }}-${{ github.event.action }}-${{ github.head_ref || github.run_id }} cancel-in-progress: true jobs: @@ -13,6 +13,7 @@ jobs: permissions: pull-requests: write contents: read + actions: read timeout-minutes: 10 steps: - uses: actions/checkout@v4 diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 3006dc7f954..c9a6e8caaf3 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -38,7 +38,7 @@ name: Docs CI jobs: typos: - if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false }} + if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1 }} name: Spell Check with Typos runs-on: ubuntu-latest steps: @@ -48,7 +48,7 @@ jobs: - uses: crate-ci/typos@v1.50.2 license-header-check: - if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false }} + if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1 }} runs-on: ubuntu-latest name: Check License Header steps: @@ -58,7 +58,7 @@ jobs: - uses: korandoru/hawkeye@v5 required-checks: - if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false }} + if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1 }} name: ${{ matrix.check }} runs-on: ubuntu-slim strategy: diff --git a/.github/workflows/grafana.yml b/.github/workflows/grafana.yml index 6f3f1aa8fdf..2ec57b3cf5d 100644 --- a/.github/workflows/grafana.yml +++ b/.github/workflows/grafana.yml @@ -2,6 +2,7 @@ name: Check Grafana Panels on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] branches: - main paths: diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index c81dba199e9..db4f4bef65b 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -31,6 +31,11 @@ name: Integration CI env: CARGO_FUZZ_VERSION: "0.13.2" + # Fuzz targets require a nightly toolchain (cargo-fuzz needs + # `-Zsanitizer=fuzzer`). This pin decouples the fuzz workflow from the + # workspace default in rust-toolchain.toml, so it keeps working when the + # workspace moves to a stable toolchain. + FUZZ_RUST_TOOLCHAIN: "nightly-2026-03-21" concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} @@ -38,7 +43,7 @@ concurrency: jobs: build: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Build GreptimeDB binaries runs-on: ${{ matrix.os }} strategy: @@ -81,7 +86,7 @@ jobs: # `cargo gc` will invoke `cargo build` with specified args run: | started_epoch="$(date +%s)" - cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions" + cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,ai_functions" elapsed_secs=$(( $(date +%s) - started_epoch )) printf -- '- Binary build: **%ss**\n' "${elapsed_secs}" >>"${GITHUB_STEP_SUMMARY}" - name: Pack greptime binaries @@ -99,7 +104,7 @@ jobs: version: current sqlness: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Sqlness Test (${{ matrix.mode.name }}) needs: build runs-on: ${{ matrix.os }} @@ -179,7 +184,7 @@ jobs: retention-days: 3 export-import-v2-e2e: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Export/Import V2 E2E Test needs: build runs-on: ubuntu-latest @@ -249,7 +254,7 @@ jobs: retention-days: 3 run-multi-lang-tests: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Run Multi-language SDK Tests needs: build uses: ./.github/workflows/run-multi-lang-tests.yml @@ -257,7 +262,7 @@ jobs: artifact-name: bins compat-updater-check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) && (github.event_name == 'merge_group' || github.event_name == 'pull_request') }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && (github.event_name == 'merge_group' || github.event_name == 'pull_request') }} name: Check compat version updater runs-on: ubuntu-latest timeout-minutes: 10 @@ -328,6 +333,7 @@ jobs: - uses: actions-rust-lang/setup-rust-toolchain@v1 with: cache: false + toolchain: ${{ env.FUZZ_RUST_TOOLCHAIN }} - name: Fuzz Dependency Cache id: fuzz-dependencies-cache uses: Swatinem/rust-cache@v2 @@ -343,7 +349,7 @@ jobs: uses: actions/cache/restore@v4 with: path: ${{ runner.temp }}/greptime-fuzz-binaries - key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', 'rust-toolchain.toml', '.cargo/**', 'tests-fuzz/**', 'src/**') }} + key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} - name: Report fuzz cache lookup shell: bash run: | @@ -355,6 +361,8 @@ jobs: - name: Set Rust Fuzz if: steps.fuzz-binaries.outputs.cache-hit != 'true' shell: bash + env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} run: | sudo apt-get install -y libfuzzer-14-dev cargo install cargo-fuzz --version "${CARGO_FUZZ_VERSION}" --locked @@ -363,6 +371,7 @@ jobs: shell: bash env: CUSTOM_LIBFUZZER_PATH: /usr/lib/llvm-14/lib/libFuzzer.a + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} FUZZ_TARGETS: | fuzz_create_database fuzz_create_table @@ -410,7 +419,7 @@ jobs: uses: actions/cache/save@v4 with: path: ${{ runner.temp }}/greptime-fuzz-binaries - key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', 'rust-toolchain.toml', '.cargo/**', 'tests-fuzz/**', 'src/**') }} + key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} - name: Pack prebuilt fuzz binaries shell: bash run: | @@ -479,6 +488,7 @@ jobs: - name: Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 FUZZ_SERVICE_LOG: /tmp/greptime-fuzz-runtime/greptime.log with: @@ -537,6 +547,7 @@ jobs: - name: Run Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 GT_FUZZ_BINARY_PATH: ./bin/greptime GT_FUZZ_INSTANCE_ROOT_DIR: /tmp/greptime-fuzz-artifacts/targets/unstable_fuzz_create_table_standalone/service/instance/ @@ -820,6 +831,7 @@ jobs: - name: Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 GT_KAFKA_WAL_HELPER_URL: http://127.0.0.1:8080 GT_KAFKA_ADDR: kafka.kafka-cluster.svc.cluster.local:9092 @@ -994,6 +1006,7 @@ jobs: - name: Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 with: targets: ${{ matrix.group.targets }} diff --git a/.github/workflows/nightly-build.yml b/.github/workflows/nightly-build.yml index 14ebb6e715e..c78ca4dc5b9 100644 --- a/.github/workflows/nightly-build.yml +++ b/.github/workflows/nightly-build.yml @@ -151,6 +151,10 @@ jobs: disable-run-tests: ${{ env.DISABLE_RUN_TESTS }} image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -173,6 +177,10 @@ jobs: disable-run-tests: ${{ env.DISABLE_RUN_TESTS }} image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index d4afa2778b8..f7f955233a1 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -71,6 +71,57 @@ jobs: path: C:\Users\RUNNER~1\AppData\Local\Temp\sqlness* retention-days: 3 + # PR CI runs the object-store matrix against a local MinIO. This job is where + # the same matrix meets real AWS S3. + test-on-linux-real-s3: + name: Run tests on Linux against real S3 + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + runs-on: ubuntu-22.04-arm + timeout-minutes: 90 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: arduino/setup-protoc@v3 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + - uses: rui314/setup-mold@v1 + - name: Install toolchain + uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + with: + shared-key: "coverage-test" + cache-all-crates: "true" + save-if: false + - name: Install latest nextest release + uses: taiki-e/install-action@nextest + - name: Setup external services + working-directory: tests-integration/fixtures + run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait + - name: Run nextest cases + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions -F hdfs-object-store + env: + CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" + RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8MB + CARGO_INCREMENTAL: 0 + GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} + GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + GT_ETCD_TLS_ENDPOINTS: https://127.0.0.1:2378 + GT_ETCD_ENDPOINTS: http://127.0.0.1:2379 + GT_POSTGRES_ENDPOINTS: postgres://greptimedb:admin@127.0.0.1:5432/postgres + GT_POSTGRES15_ENDPOINTS: postgres://test_user:test_password@127.0.0.1:5433/postgres + GT_POSTGRES15_SCHEMA: test_schema + GT_MYSQL_ENDPOINTS: mysql://greptimedb:admin@127.0.0.1:3306/mysql + GT_KAFKA_ENDPOINTS: 127.0.0.1:9092 + GT_KAFKA_SASL_ENDPOINTS: 127.0.0.1:9093 + UNITTEST_LOG_DIR: "__unittest_logs" + test-on-windows: name: Run tests on Windows if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} @@ -130,7 +181,7 @@ jobs: check-status: name: Check status - needs: [sqlness-test, sqlness-windows, test-on-windows] + needs: [sqlness-test, sqlness-windows, test-on-windows, test-on-linux-real-s3] if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest outputs: diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 8c9a93789cf..6a368c455d6 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -190,6 +190,12 @@ jobs: CARGO_HOME: /home/runner/.cargo UV_CACHE_DIR: /home/runner/.cargo/uv-cache RUSTUP_HOME: /opt/rustup + # Deliberately NOT derived from rust-toolchain.toml: the benchmark + # also compiles the BASE checkout (the previous nightly build, which + # may predate the stable-toolchain migration and still uses + # #![feature] gates), so the runner must keep a nightly toolchain + # that can build historical revisions. This pin moves only via a + # deliberate runner-image rebuild. RUSTUP_TOOLCHAIN: nightly-2026-03-21 RUSTUP_AUTO_INSTALL: "0" CARGO_TARGET_DIR: /home/runner/query-regression-target @@ -504,6 +510,9 @@ jobs: require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo" require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc" require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)' + # The runner toolchain is pinned to nightly-2026-03-21 (see the + # RUSTUP_TOOLCHAIN comment): it must compile historical base + # checkouts, so it intentionally does not follow rust-toolchain.toml. require_match cargo "$(capture cargo --version)" \ '^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' require_match rustc "$(capture rustc --version)" \ @@ -540,7 +549,7 @@ jobs: readonly EXPECTED_SCCACHE_DIR="/home/runner/.cache/sccache" readonly EXPECTED_RUSTC_WRAPPER="/usr/local/bin/sccache" readonly RUNNER_IMAGE_DIGEST="sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11" - readonly RUNNER_IMAGE_EPOCH="5" + readonly RUNNER_IMAGE_EPOCH="7" require_expected_root() { local name="$1" diff --git a/.github/workflows/release-dev-builder-images.yaml b/.github/workflows/release-dev-builder-images.yaml index bebd5ed61e7..2e1041c1ff4 100644 --- a/.github/workflows/release-dev-builder-images.yaml +++ b/.github/workflows/release-dev-builder-images.yaml @@ -1,4 +1,4 @@ -name: Release dev-builder images +name: Release builder images on: push: @@ -7,6 +7,7 @@ on: paths: - rust-toolchain.toml - 'docker/dev-builder/**' + - '.github/runner-scale-sets/query-regression/**' workflow_dispatch: # Allows you to run this workflow manually. inputs: release_dev_builder_ubuntu_image: @@ -29,6 +30,11 @@ on: description: Release dev-builder-riscv64 image required: false default: false + release_query_regression_runner_image: + type: boolean + description: Rebuild the query-regression ECS runner image (and update QUERY_REGRESSION_ECS_IMAGE_ID + RUNNER_IMAGE_EPOCH) + required: false + default: false update_dev_builder_image_tag: type: boolean description: Update the DEV_BUILDER_IMAGE_TAG in Makefile and create a PR @@ -36,12 +42,59 @@ on: default: false jobs: + changes: + name: Detect changed components + runs-on: ubuntu-latest + outputs: + dev-builder: ${{ steps.filter.outputs.dev-builder }} + query-regression-runner: ${{ steps.filter.outputs.query-regression-runner }} + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + persist-credentials: false + - name: Filter changes + id: filter + shell: bash + run: | + set -euo pipefail + base="${{ github.event.before }}" + head="${{ github.sha }}" + if [[ "${{ github.event_name }}" != "push" ]] || [[ -z "${base}" ]] || ! git cat-file -e "${base}^{commit}" 2>/dev/null; then + # Dispatch, or an unknown base (new branch/force-push): treat + # everything as changed so gating falls back to job inputs. + echo "dev-builder=true" >> "$GITHUB_OUTPUT" + echo "query-regression-runner=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + files="$(git diff --name-only "${base}" "${head}")" + dev_builder=false + query_regression_runner=false + # rust-toolchain.toml only gates the dev-builder images: they bake + # the workspace toolchain. The query-regression runner toolchain is + # pinned inside its Dockerfile (nightly-2026-03-21) and must NOT + # follow the workspace pin, so its rebuild triggers only on changes + # under its own directory. + if grep -qx 'rust-toolchain.toml' <<<"${files}"; then + dev_builder=true + fi + if grep -q '^docker/dev-builder/' <<<"${files}"; then + dev_builder=true + fi + if grep -q '^\.github/runner-scale-sets/query-regression/' <<<"${files}"; then + query_regression_runner=true + fi + echo "dev-builder=${dev_builder}" >> "$GITHUB_OUTPUT" + echo "query-regression-runner=${query_regression_runner}" >> "$GITHUB_OUTPUT" + release-dev-builder-images: name: Release dev builder images # The jobs are triggered by the following events: # 1. Manually triggered workflow_dispatch event # 2. Push event when the PR that modifies the `rust-toolchain.toml` or `docker/dev-builder/**` is merged to main - if: ${{ github.event_name == 'push' || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image || inputs.release_dev_builder_riscv64_image }} + needs: [changes] + if: ${{ (github.event_name == 'push' && needs.changes.outputs.dev-builder == 'true') || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image || inputs.release_dev_builder_riscv64_image }} runs-on: ubuntu-latest outputs: version: ${{ steps.set-version.outputs.version }} @@ -59,7 +112,10 @@ jobs: commitShortSHA=`echo ${{ github.sha }} | cut -c1-8` buildTime=`date +%Y%m%d%H%M%S` BUILD_VERSION="$commitShortSHA-$buildTime" - RUST_TOOLCHAIN_VERSION=$(cat rust-toolchain.toml | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') + # Extract the full channel string (e.g. "1.96.1" or + # "nightly-2026-03-21") so the image version always carries the + # toolchain it was built with. + RUST_TOOLCHAIN_VERSION=$(grep -E '^channel' rust-toolchain.toml | cut -d'"' -f2) IMAGE_VERSION="${RUST_TOOLCHAIN_VERSION}-${BUILD_VERSION}" echo "VERSION=${IMAGE_VERSION}" >> $GITHUB_ENV echo "version=$IMAGE_VERSION" >> $GITHUB_OUTPUT @@ -278,3 +334,120 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | ./.github/scripts/update-dev-builder-version.sh ${{ needs.release-dev-builder-images.outputs.version }} + + rebuild-query-regression-runner-image: + # Rebuilds the query-regression ECS runner image via the ops tool in + # .github/runner-scale-sets/query-regression/ecs-image/: boots a + # temporary pay-as-you-go ECS builder, snapshots a new custom image, + # then completes the documented lockstep updates: + # 1. bumps RUNNER_IMAGE_EPOCH in query-regression.yml (target-cache + # invalidation) via a PR from a timestamped ci/ branch (main is + # branch-protected; mirrors update-dev-builder-version.sh), and + # 2. points the repo variable QUERY_REGRESSION_ECS_IMAGE_ID at the + # new image. The runner toolchain is pinned inside the Dockerfile + # (nightly-2026-03-21, independent of rust-toolchain.toml — the + # benchmark must compile historical base checkouts), so image + # rebuilds do not change the cache ABI. + # + # Required repository configuration (same names the provisioning job + # uses): vars ALIYUN_ECS_REGION_ID, ALIYUN_ECS_VSWITCH_ID, + # ALIYUN_ECS_SECURITY_GROUP_ID, optionally ALIYUN_ECS_RESOURCE_GROUP_ID + # and ALIYUN_ECS_BASE_IMAGE_ID (deterministic base-image pin; otherwise + # the rebuild auto-resolves the latest public Ubuntu 26.04 image). + # Secrets: ALICLOUD_ECS_ACCESS_KEY_ID, ALICLOUD_ECS_ACCESS_KEY_SECRET, + # GH_PERSONAL_ACCESS_TOKEN (repo push + actions-variable write). + name: Rebuild query-regression runner image + needs: [changes] + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'push' && needs.changes.outputs.query-regression-runner == 'true') || inputs.release_query_regression_runner_image) }} + runs-on: ubuntu-latest + timeout-minutes: 150 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + # Keep the PAT credential for the epoch-bump push below. + persist-credentials: true + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Rebuild the runner image + id: rebuild + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + run: | + set -euo pipefail + args=( + --region-id "${{ vars.ALIYUN_ECS_REGION_ID }}" + --vswitch-id "${{ vars.ALIYUN_ECS_VSWITCH_ID }}" + --security-group-id "${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }}" + --base-image-id "${{ vars.ALIYUN_ECS_BASE_IMAGE_ID }}" + ) + if [[ -n "${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }}" ]]; then + args+=(--resource-group-id "${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }}") + fi + uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + "${args[@]}" 2>&1 | tee rebuild.log + image_id="$(sed -n 's/^Custom image ready: \([^ ]*\) .*/\1/p' rebuild.log | tail -n1)" + if [[ -z "${image_id}" ]]; then + echo "Error: no image id found in build output" >&2 + exit 1 + fi + echo "image_id=${image_id}" >> "$GITHUB_OUTPUT" + echo "## Rebuilt runner image" >> "$GITHUB_STEP_SUMMARY" + echo "New image: \`${image_id}\`" >> "$GITHUB_STEP_SUMMARY" + + - name: Open the epoch-bump PR and update the image id variable + env: + GH_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + run: | + set -euo pipefail + image_id="${{ steps.rebuild.outputs.image_id }}" + old_epoch="$(grep -oE 'readonly RUNNER_IMAGE_EPOCH="[0-9]+"' \ + .github/workflows/query-regression.yml | grep -oE '[0-9]+' | head -n1)" + if [[ -z "${old_epoch}" ]]; then + echo "Error: RUNNER_IMAGE_EPOCH not found in query-regression.yml" >&2 + exit 1 + fi + new_epoch=$((old_epoch + 1)) + + # main is branch-protected (required reviews + status checks), so + # the epoch bump cannot be pushed there directly. Mirror + # .github/scripts/update-dev-builder-version.sh: timestamped bot + # branch, plain push, and a PR with reviewers. Do NOT use [skip ci] + # on the commit: the required checks must be able to run for the PR + # to become mergeable. + BRANCH="ci/update-query-regression-epoch-$(date +%Y%m%d%H%M%S)" + git config user.name "greptimedb-ci" + git config user.email "greptimedb-ci@greptime.com" + git fetch origin main + git checkout -b "${BRANCH}" origin/main + sed -i "s/readonly RUNNER_IMAGE_EPOCH=\"${old_epoch}\"/readonly RUNNER_IMAGE_EPOCH=\"${new_epoch}\"/" \ + .github/workflows/query-regression.yml + git add .github/workflows/query-regression.yml + git commit -s -m "ci(query-regression): bump RUNNER_IMAGE_EPOCH to ${new_epoch} for image ${image_id}" + git push origin "${BRANCH}" + + gh pr create \ + --title "ci(query-regression): bump RUNNER_IMAGE_EPOCH to ${new_epoch}" \ + --body "Auto-generated by the \`Rebuild query-regression runner image\` job in ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID} for image \`${image_id}\`. Merging this invalidates the query-regression target cache (RUNNER_IMAGE_EPOCH ${old_epoch} → ${new_epoch}). The \`QUERY_REGRESSION_ECS_IMAGE_ID\` repo variable already points at the new image; the epoch bump is belt-and-suspenders for image-content drift." \ + --base main \ + --head "${BRANCH}" \ + --reviewer sunng87 \ + --reviewer daviderli614 \ + --reviewer killme2008 \ + --reviewer evenyag \ + --reviewer fengjiachun \ + --reviewer WenyXu + + # Point the variable at the new image once the PR exists. Image + # rebuilds keep the same pinned toolchain, so the cache ABI is + # unchanged and not waiting for the PR merge is fine. + gh api -X PATCH "repos/${{ github.repository }}/actions/variables/QUERY_REGRESSION_ECS_IMAGE_ID" \ + -f value="${image_id}" + + echo "## Lockstep updates" >> "$GITHUB_STEP_SUMMARY" + echo "- QUERY_REGRESSION_ECS_IMAGE_ID → \`${image_id}\`" >> "$GITHUB_STEP_SUMMARY" + echo "- RUNNER_IMAGE_EPOCH → \`${new_epoch}\` (PR on \`${BRANCH}\`, pending review/merge)" >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index acf52e241c2..0fb22f23e61 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -4,15 +4,20 @@ name: Release # 1. The tag('v*.*.*') push release: the release workflow will be triggered by the tag push event. # 2. The scheduled release(the version will be '-nightly-YYYYMMDD', any pre-release # extension in the next release version like '-alpha.1' is stripped, so 'nightly' is the extension): -# the release workflow will be triggered by the schedule event. +# the release workflow will be triggered by the schedule event. The schedule runs every weekday, +# but the release only proceeds when the latest published nightly release is older than +# NIGHTLY_RELEASE_MAX_AGE_DAYS days, so a failed nightly run is retried by the next weekday +# schedule instead of leaving a whole week without nightly builds. # A failed tag push release can be manually dispatched on the same tag to recover it without changing the formal release version. on: push: tags: - "v*.*.*" schedule: - # At 00:00 on Monday. - - cron: '0 0 * * 1' + # At 00:00 on every weekday (Monday to Friday). The scheduled nightly release only + # proceeds when the latest published nightly release is older than + # NIGHTLY_RELEASE_MAX_AGE_DAYS days, see the 'Check nightly release necessity' step. + - cron: '0 0 * * 1-5' workflow_dispatch: # Allows you to run this workflow manually. # Notes: The GitHub Actions ONLY support 10 inputs; this final input is the tenth. inputs: @@ -79,7 +84,7 @@ on: default: false release_validation: type: choice - description: Manual policy for accepted perf regressions or runner failures; automatic releases always run all validation + description: Manual policy for accepted perf regressions or runner failures; automatic releases always run all validation (scheduled nightlies run query regression non-blocking) required: true default: all options: @@ -100,6 +105,10 @@ env: # the pre-release extension of the next release version is stripped, e.g. v0.2.0-alpha.1 -> v0.2.0-nightly-20230313. NIGHTLY_RELEASE_PREFIX: nightly + # The maximum age (in days) of the latest published nightly release before a scheduled + # nightly release may run again; see the 'Check nightly release necessity' step. + NIGHTLY_RELEASE_MAX_AGE_DAYS: 5 + jobs: allocate-runners: name: Allocate runners @@ -123,6 +132,10 @@ jobs: # The 'is-current-version-latest' determines whether to update 'latest' Docker tags and downstream repositories. is-current-version-latest: ${{ steps.check-version.outputs.is-current-version-latest }} is-current-version-stable: ${{ steps.check-version.outputs.is-current-version-stable }} + + # Whether a scheduled nightly release should run: it only runs when the latest published + # nightly release is older than NIGHTLY_RELEASE_MAX_AGE_DAYS days. Always 'true' for other events. + nightly-required: ${{ steps.check-nightly.outputs.nightly-required }} steps: - name: Checkout uses: actions/checkout@v4 @@ -135,6 +148,19 @@ jobs: run: | ./scripts/check-builder-rust-version.sh + # A scheduled run only proceeds when the latest published nightly release is older than + # NIGHTLY_RELEASE_MAX_AGE_DAYS days, so a failed nightly run is retried by the next weekday + # schedule. The check runs before the runners are allocated to avoid wasting EC2 resources. + - name: Check nightly release necessity + id: check-nightly + run: | + echo "nightly-required=$(./.github/scripts/check-nightly-release.sh)" >> "$GITHUB_OUTPUT" + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_EVENT_NAME: ${{ github.event_name }} + GITHUB_REPOSITORY: ${{ github.repository }} + NIGHTLY_RELEASE_MAX_AGE_DAYS: ${{ env.NIGHTLY_RELEASE_MAX_AGE_DAYS }} + # The create-version will create a global variable named 'version' in the global workflows. # - If it's a tag push release or manual dispatch on a tag with REUSE_EXISTING_RELEASE_TAG=true, the version is the tag name(${{ github.ref_name }}); # - If it's a scheduled release, the version is '-nightly-$buildTime', like v0.2.0-nightly-20230313 (pre-release extensions like '-alpha.1' are stripped); @@ -156,7 +182,7 @@ jobs: ./.github/scripts/check-version.sh "${{ steps.create-version.outputs.version }}" - name: Allocate linux-amd64 runner - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-amd64-runner with: @@ -170,7 +196,7 @@ jobs: subnet-id: ${{ vars.EC2_RUNNER_SUBNET_ID }} - name: Allocate linux-arm64 runner - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-arm64-runner with: @@ -185,7 +211,7 @@ jobs: build-linux-amd64-artifacts: name: Build linux-amd64 artifacts - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -204,10 +230,14 @@ jobs: disable-run-tests: ${{ env.DISABLE_RUN_TESTS }} image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -227,10 +257,14 @@ jobs: image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: true + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-riscv64-artifacts: name: Build linux-riscv64 artifacts - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -294,7 +328,7 @@ jobs: needs: [ allocate-runners, ] - if: ${{ inputs.build_macos_artifacts || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_macos_artifacts || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} steps: - uses: actions/checkout@v4 with: @@ -332,7 +366,7 @@ jobs: needs: [ allocate-runners, ] - if: ${{ inputs.build_windows_artifacts || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_windows_artifacts || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} steps: - run: git config --global core.autocrlf false @@ -357,7 +391,7 @@ jobs: prepare-release-validation: name: Prepare release validation - if: ${{ inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -434,6 +468,9 @@ jobs: if [[ "${validation_policy}" != "all" ]]; then echo "- ⚠️ Validation bypassed: compatibility=${run_compat}, query-regression=${run_query_regression}" fi + if [[ "${GITHUB_EVENT_NAME}" == "schedule" ]]; then + echo "- ℹ️ Query-regression failure or cancellation is non-blocking for scheduled nightly releases" + fi } >> "${GITHUB_STEP_SUMMARY}" compat-release: @@ -517,9 +554,13 @@ jobs: release-images-to-dockerhub: name: Build and push images to DockerHub + # Scheduled nightly releases do not gate on query regression: the test + # still runs and reports, but a failure or cancellation does not block + # the nightly release. if: | always() && (inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && + (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') && needs.allocate-runners.result == 'success' && needs.build-linux-amd64-artifacts.result == 'success' && needs.build-linux-arm64-artifacts.result == 'success' && @@ -527,7 +568,8 @@ jobs: (needs.compat-release.result == 'success' || (needs.prepare-release-validation.outputs.run-compat == 'false' && needs.compat-release.result == 'skipped')) && (needs.query-regression-release.result == 'success' || - (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped')) + (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped') || + github.event_name == 'schedule') needs: [ allocate-runners, build-linux-amd64-artifacts, @@ -563,7 +605,18 @@ jobs: release-cn-artifacts: name: Release artifacts to CN region - if: ${{ inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + # A non-blocking query regression failure must not skip downstream releases. + if: | + always() && + (inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && + (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') && + needs.allocate-runners.result == 'success' && + needs.build-linux-amd64-artifacts.result == 'success' && + needs.build-linux-arm64-artifacts.result == 'success' && + needs.build-linux-riscv64-artifacts.result == 'success' && + needs.build-macos-artifacts.result == 'success' && + needs.build-windows-artifacts.result == 'success' && + needs.release-images-to-dockerhub.result == 'success' needs: [ # The job have to wait for all the artifacts are built. allocate-runners, build-linux-amd64-artifacts, @@ -607,15 +660,19 @@ jobs: name: Create GitHub release and upload artifacts # Use always() to run even when optional jobs (macos, windows) are skipped. # Then check that required jobs succeeded and optional jobs didn't fail. + # Scheduled nightly releases do not gate on query regression either + # (see release-images-to-dockerhub); the test still runs and reports. if: | always() && (inputs.publish_github_release || github.event_name == 'push' || github.event_name == 'schedule') && + (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') && needs.allocate-runners.result == 'success' && needs.prepare-release-validation.result == 'success' && (needs.compat-release.result == 'success' || (needs.prepare-release-validation.outputs.run-compat == 'false' && needs.compat-release.result == 'skipped')) && (needs.query-regression-release.result == 'success' || - (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped')) && + (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped') || + github.event_name == 'schedule') && needs.build-linux-amd64-artifacts.result == 'success' && needs.build-linux-arm64-artifacts.result == 'success' && needs.build-linux-riscv64-artifacts.result == 'success' && @@ -718,7 +775,21 @@ jobs: bump-downstream-repo-versions: name: Bump downstream repo versions - if: ${{ github.event_name == 'schedule' || ((github.event_name == 'push' || github.event_name == 'workflow_dispatch') && github.ref_type == 'tag' && !contains(github.ref_name, 'nightly')) && (needs.allocate-runners.outputs.is-current-version-stable == 'true' && needs.allocate-runners.outputs.is-current-version-latest == 'true') }} + # A non-blocking query regression failure must not skip downstream updates. + if: | + always() && + needs.allocate-runners.result == 'success' && + needs.publish-github-release.result == 'success' && + ( + (github.event_name == 'schedule' && needs.allocate-runners.outputs.nightly-required == 'true') || + ( + (github.event_name == 'push' || github.event_name == 'workflow_dispatch') && + github.ref_type == 'tag' && + !contains(github.ref_name, 'nightly') && + needs.allocate-runners.outputs.is-current-version-stable == 'true' && + needs.allocate-runners.outputs.is-current-version-latest == 'true' + ) + ) needs: [allocate-runners, publish-github-release] runs-on: ubuntu-latest # Permission reference: https://docs.github.com/en/actions/using-jobs/assigning-permissions-to-jobs @@ -785,9 +856,10 @@ jobs: ./.github/scripts/update-homebrew-greptme-version.sh notification: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'push' || github.event_name == 'schedule') && always() }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'push' || (github.event_name == 'schedule' && needs.allocate-runners.outputs.nightly-required == 'true')) && always() }} name: Send notification to Greptime team needs: [ + allocate-runners, release-images-to-dockerhub, build-macos-artifacts, build-windows-artifacts, diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 1a093a2d2a5..e29b69678c1 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -24,7 +24,7 @@ on: name: Rust CI env: - CARGO_UDEPS_VERSION: "0.1.61" + CARGO_SHEAR_VERSION: "1.13.4" concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} @@ -32,7 +32,7 @@ concurrency: jobs: fmt: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Rustfmt runs-on: ubuntu-latest timeout-minutes: 60 @@ -50,7 +50,7 @@ jobs: run: make fmt-check clippy: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Clippy runs-on: ubuntu-latest timeout-minutes: 60 @@ -76,7 +76,7 @@ jobs: run: make clippy check-udeps: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check Unused Dependencies runs-on: ubuntu-latest timeout-minutes: 60 @@ -84,26 +84,17 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 - with: - repo-token: ${{ secrets.GITHUB_TOKEN }} - uses: actions-rust-lang/setup-rust-toolchain@v1 - - name: Rust Cache - uses: Swatinem/rust-cache@v2 - with: - shared-key: "check-udeps" - cache-all-crates: "true" - save-if: ${{ github.ref == 'refs/heads/main' }} - - name: Install cargo-udeps + - name: Install cargo-shear uses: taiki-e/install-action@v2 with: - tool: cargo-udeps@${{ env.CARGO_UDEPS_VERSION }} + tool: cargo-shear@${{ env.CARGO_SHEAR_VERSION }} fallback: none - name: Check unused dependencies run: make check-udeps check-riscv64: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check riscv64 build runs-on: ubuntu-latest timeout-minutes: 90 @@ -129,7 +120,7 @@ jobs: run: cargo check --locked --workspace --all-targets --features servers/dashboard --target riscv64gc-unknown-linux-gnu check-windows: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check Windows build runs-on: windows-2022 timeout-minutes: 60 @@ -157,8 +148,50 @@ jobs: env: RUSTFLAGS: "-D warnings" - conflict-check: + test-windows: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + name: Test Windows + runs-on: windows-2022 + timeout-minutes: 90 + env: + CARGO_INCREMENTAL: 0 + CARGO_PROFILE_DEV_DEBUG: 0 + CARGO_PROFILE_TEST_DEBUG: 0 + RUSTFLAGS: "-D warnings -C linker=lld-link" + RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8MB + 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 + with: + cache: false + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + with: + shared-key: "test-windows" + cache-all-crates: "true" + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Install Cargo Nextest + uses: taiki-e/install-action@nextest + - name: Install WSL distribution + uses: Vampire/setup-wsl@v2 + with: + distribution: Ubuntu-22.04 + - name: Run Windows tests + run: cargo nextest run --locked --workspace -F dashboard --no-fail-fast + + conflict-check: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check for conflict runs-on: ubuntu-latest steps: @@ -169,7 +202,7 @@ jobs: uses: olivernybroe/action-conflict-finder@v4.0 test: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && github.event_name != 'merge_group' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && github.event_name != 'merge_group' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} runs-on: ubuntu-22.04-arm timeout-minutes: 60 needs: [conflict-check, clippy, fmt, check-udeps] @@ -200,21 +233,20 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions -F hdfs-object-store env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 RUST_MIN_STACK: 8388608 # 8MB CARGO_INCREMENTAL: 0 - GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} - GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} - GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - GT_MINIO_BUCKET: greptime - GT_MINIO_ACCESS_KEY_ID: superpower_ci_user - GT_MINIO_ACCESS_KEY: superpower_password - GT_MINIO_REGION: us-west-2 - GT_MINIO_ENDPOINT_URL: http://127.0.0.1:9000 + # Runs the object-store matrix against the local MinIO from + # tests-integration/fixtures. Real AWS S3 is covered by the nightly + # jobs and by the release build. + GT_S3_BUCKET: greptime + GT_S3_ACCESS_KEY_ID: superpower_ci_user + GT_S3_ACCESS_KEY: superpower_password + GT_S3_REGION: us-west-2 + GT_S3_ENDPOINT_URL: http://127.0.0.1:9000 GT_ETCD_TLS_ENDPOINTS: https://127.0.0.1:2378 GT_ETCD_ENDPOINTS: http://127.0.0.1:2379 GT_POSTGRES_ENDPOINTS: postgres://greptimedb:admin@127.0.0.1:5432/postgres @@ -258,20 +290,19 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions + run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions -F hdfs-object-store env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 CARGO_INCREMENTAL: 0 - GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} - GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} - GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - GT_MINIO_BUCKET: greptime - GT_MINIO_ACCESS_KEY_ID: superpower_ci_user - GT_MINIO_ACCESS_KEY: superpower_password - GT_MINIO_REGION: us-west-2 - GT_MINIO_ENDPOINT_URL: http://127.0.0.1:9000 + # Runs the object-store matrix against the local MinIO from + # tests-integration/fixtures. Real AWS S3 is covered by the nightly + # jobs and by the release build. + GT_S3_BUCKET: greptime + GT_S3_ACCESS_KEY_ID: superpower_ci_user + GT_S3_ACCESS_KEY: superpower_password + GT_S3_REGION: us-west-2 + GT_S3_ENDPOINT_URL: http://127.0.0.1:9000 GT_ETCD_TLS_ENDPOINTS: https://127.0.0.1:2378 GT_ETCD_ENDPOINTS: http://127.0.0.1:2379 GT_POSTGRES_ENDPOINTS: postgres://greptimedb:admin@127.0.0.1:5432/postgres diff --git a/.github/workflows/slash-command-dispatch.yml b/.github/workflows/slash-command-dispatch.yml index 26fa27b79dd..c47a9ed9ced 100644 --- a/.github/workflows/slash-command-dispatch.yml +++ b/.github/workflows/slash-command-dispatch.yml @@ -1,13 +1,13 @@ name: Slash Command Dispatch -# ChatOps front door: parse `/command` on PR comments, check admin -# permission, and repository_dispatch to a per-command handler. Handlers +# ChatOps front door: parse `/command` on PR comments, apply command-specific +# permissions, and repository_dispatch to a per-command handler. Handlers # own allowlists, SHA admission, and the actual work. Same-repo dispatch # uses github.token with contents: write; GitHub starts the handler run # for GITHUB_TOKEN-created repository_dispatch events. Do not pass the # long-lived GH_PERSONAL_ACCESS_TOKEN into this third-party action. # -# To add a command: list it under `commands` and add a workflow with +# To add a command: add it to `config` and add a workflow with # `on.repository_dispatch.types: ["-command"]`. on: @@ -30,8 +30,17 @@ jobs: uses: peter-evans/slash-command-dispatch@9bdcd7914ec1b75590b790b844aa3b8eee7c683a # v5.0.2 with: token: ${{ github.token }} - permission: admin - issue-type: pull-request - commands: | - query-regression - ci + # CI admission re-fetches the comment and checks author/member permissions. + config: >- + [ + {"command": "query-regression", "permission": "admin", "issue_type": "pull-request"}, + {"command": "ci", "permission": "none", "issue_type": "pull-request"} + ] + - name: Report CI dispatch failure + if: ${{ failure() && (github.event.comment.body == '/ci' || startsWith(github.event.comment.body, '/ci ')) }} + env: + GH_TOKEN: ${{ github.token }} + PR_NUMBER: ${{ github.event.issue.number }} + run: | + gh api --method POST "/repos/${GITHUB_REPOSITORY}/issues/${PR_NUMBER}/comments" \ + -f "body=CI command dispatch failed. See ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID} for details." diff --git a/.github/workflows/tracesbench.yml b/.github/workflows/tracesbench.yml new file mode 100644 index 00000000000..049faaef1a7 --- /dev/null +++ b/.github/workflows/tracesbench.yml @@ -0,0 +1,495 @@ +name: Tracesbench +on: + workflow_dispatch: + inputs: + run_greptimedb: + description: Run GreptimeDB + type: boolean + default: true + required: false + run_victoriatraces: + description: Run VictoriaTraces + type: boolean + default: false + required: false + run_tempo: + description: Run Tempo + type: boolean + default: false + required: false + compare_greptimedb: + description: Compare baseline and candidate sequentially on the same host (adds one full GreptimeDB run) + type: boolean + default: false + required: false + baseline_greptimedb_tag: + description: Baseline Docker Hub greptime/greptimedb tag; required when comparison is enabled + type: string + default: v1.2.1 + required: false + greptimedb_tag: + description: Candidate Docker Hub greptime/greptimedb tag + type: string + default: latest + required: true + dataset: + description: Dataset profile (main is experimental; native queries may exceed memory limits) + type: choice + options: [sanity, preflight, main] + default: sanity + required: true + runtime_image: + description: 'Runtime image: auto = AWS ECR Public or Aliyun ACR by provider; explicit tag/digest overrides' + type: string + default: auto + required: true + provider: + description: Cloud provider for the benchmark runner + type: choice + options: [Aliyun, AWS] + default: Aliyun + required: true + instance_type: + description: 'Instance type; auto = Aliyun: ecs.g9i.2xlarge, AWS: m7i.2xlarge (both 8 vCPU / 32 GiB). Enter a specific type to override.' + type: string + default: auto + required: true + benchmark_root: + description: Absolute disk-backed work directory (outside the o11ybench checkout) + type: string + default: /home/runner/benchmark-data + required: true + system_disk_gib: + description: 'System disk GiB: auto (or 0) = 40 for sanity/preflight, 200 for main; enter 20..2048 to override' + type: string + default: auto + required: true + benchmark_timeout_minutes: + description: Benchmark job timeout in minutes (1..360) + type: number + default: 360 + required: true + janitor_ttl_hours: + description: Instance lifetime before janitor eligibility (1..168 hours; at least timeout + 2h) + type: number + default: 8 + required: true + db_cpus: + description: CPU limit per DB container (default 8 vCPU) + type: string + default: '8' + required: true + db_memory: + description: Memory limit per DB container (default 32 GiB; shared host also runs the runtime and runner) + type: string + default: 32g + required: true + load_workers: + description: CPU workers for parallel import reading/parsing across all engines (separate from write concurrency) + type: number + default: 4 + required: true + query_timeout_seconds: + description: Timeout per query request in seconds (correctness, warmup and measurement; 1..3600) + type: number + default: 300 + required: true + runtime_cpus: + description: CPU limit per runtime container (0 = unrestricted) + type: string + default: '0' + required: true + runtime_memory: + description: Memory limit per runtime container (0 = unrestricted; separate from DB memory) + type: string + default: '0' + required: true +permissions: + contents: read +env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + BASELINE_GREPTIMEDB_TAG: ${{ inputs.baseline_greptimedb_tag }} + DATASET: ${{ inputs.dataset }} + GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} + RUNTIME_CPUS: ${{ inputs.runtime_cpus }} + RUNTIME_MEMORY: ${{ inputs.runtime_memory }} + BENCHMARK_ROOT: ${{ inputs.benchmark_root }} + RUNTIME_IMAGE: ${{ inputs.runtime_image != 'auto' && inputs.runtime_image || inputs.provider == 'AWS' && 'public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime:77a84e1c8136@sha256:c482993ec177b026235d6cd0c61d4fb6d6f1088bc78bdc0e5964a5398de6dba0' || 'greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:77a84e1c8136@sha256:c482993ec177b026235d6cd0c61d4fb6d6f1088bc78bdc0e5964a5398de6dba0' }} + PROVIDER: ${{ inputs.provider }} + REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} + SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} + BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} + JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + DB_CPUS: ${{ inputs.db_cpus }} + DB_MEMORY: ${{ inputs.db_memory }} + QUERY_TIMEOUT_SECONDS: ${{ inputs.query_timeout_seconds }} + LOAD_WORKERS: ${{ inputs.load_workers }} +jobs: + validate: + runs-on: ubuntu-latest + timeout-minutes: 15 + outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} + targets: ${{ steps.inputs.outputs.targets }} + greptimedb: ${{ steps.inputs.outputs.greptimedb }} + victoriatraces: ${{ steps.inputs.outputs.victoriatraces }} + tempo: ${{ steps.inputs.outputs.tempo }} + system_disk_gib: ${{ steps.runtime.outputs.system_disk_gib }} + runtime_image: ${{ steps.runtime.outputs.runtime_image }} + o11ybench_ref: ${{ steps.runtime.outputs.o11ybench_ref }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Validate selected targets + id: inputs + env: + greptimedb: ${{ inputs.run_greptimedb }} + victoriatraces: ${{ inputs.run_victoriatraces }} + tempo: ${{ inputs.run_tempo }} + run: | + set -euo pipefail + selected= + for target in greptimedb victoriatraces tempo; do + enabled=${!target} + [[ "$enabled" == true || "$enabled" == false ]] + printf '%s=%s\n' "$target" "$enabled" >> "$GITHUB_OUTPUT" + if [[ "$enabled" == true ]]; then selected="${selected:+$selected,}$target"; fi + done + [[ -n "$selected" ]] || { echo 'Select at least one database' >&2; exit 1; } + if [[ "$greptimedb" == true ]]; then + [[ "$GREPTIMEDB_TAG" =~ ^[a-zA-Z0-9_][a-zA-Z0-9_.-]{0,127}$ ]] || { echo 'Invalid GreptimeDB image tag' >&2; exit 1; } + fi + printf 'targets=%s\n' "$selected" >> "$GITHUB_OUTPUT" + - name: Validate version comparison + env: + RUN_GREPTIMEDB: ${{ inputs.run_greptimedb }} + run: | + python3 greptimedb/tests/perf/test_benchmark_version_comparison.py + python3 greptimedb/.github/scripts/benchmark-version-comparison.py validate + - name: Resolve CI runner + id: runner + run: | + set -euo pipefail + instance_type="$REQUESTED_INSTANCE_TYPE" + if [[ "$instance_type" == auto ]]; then + case "$PROVIDER" in + Aliyun) instance_type=ecs.g9i.2xlarge ;; + AWS) instance_type=m7i.2xlarge ;; + esac + fi + python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$instance_type" + - name: Validate budgets and shared runtime + id: runtime + run: | + set -euo pipefail + [[ "$BENCHMARK_ROOT" == /* && "$BENCHMARK_ROOT" != / && "$BENCHMARK_ROOT" != */ ]] + [[ "$BENCHMARK_ROOT" != /tmp && "$BENCHMARK_ROOT" != /tmp/* && "$BENCHMARK_ROOT" != *,* ]] + [[ "$BENCHMARK_ROOT" != *$'\n'* && "$BENCHMARK_ROOT" != *$'\r'* ]] + [[ "/${BENCHMARK_ROOT#/}/" != */../* && "/${BENCHMARK_ROOT#/}/" != */./* ]] + case "$DATASET" in sanity|preflight|main) ;; *) echo "Unknown dataset"; exit 1;; esac + if [[ "$SYSTEM_DISK_GIB" == auto || "$SYSTEM_DISK_GIB" == 0 ]]; then + if [[ "$DATASET" == main ]]; then SYSTEM_DISK_GIB=200; else SYSTEM_DISK_GIB=40; fi + fi + for value in "$SYSTEM_DISK_GIB" "$BENCHMARK_TIMEOUT_MINUTES" "$JANITOR_TTL_HOURS"; do + [[ "$value" =~ ^[1-9][0-9]{0,3}$ ]] || { echo 'Resource budgets must be positive integers'; exit 1; } + done + ((SYSTEM_DISK_GIB >= 20 && SYSTEM_DISK_GIB <= 2048)) + ((BENCHMARK_TIMEOUT_MINUTES <= 360 && JANITOR_TTL_HOURS <= 168)) + ((JANITOR_TTL_HOURS * 60 >= BENCHMARK_TIMEOUT_MINUTES + 120)) + [[ "$DB_CPUS" =~ ^[1-9][0-9]*$ && "$DB_MEMORY" =~ ^[1-9][0-9]*[mMgG]$ ]] + [[ "$RUNTIME_CPUS" =~ ^(0|[1-9][0-9]*)$ && "$RUNTIME_MEMORY" =~ ^(0|[1-9][0-9]*[mMgG])$ ]] + [[ "$QUERY_TIMEOUT_SECONDS" =~ ^[1-9][0-9]{0,3}$ ]] + ((QUERY_TIMEOUT_SECONDS <= 3600)) + [[ "$LOAD_WORKERS" =~ ^[1-9][0-9]{0,2}$ ]] + ((LOAD_WORKERS <= 256)) + runtime_limits=() + [[ "$RUNTIME_CPUS" == 0 ]] || runtime_limits+=(--cpus "$RUNTIME_CPUS") + [[ "$RUNTIME_MEMORY" == 0 ]] || runtime_limits+=(--memory "$RUNTIME_MEMORY" --memory-swap "$RUNTIME_MEMORY") + [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* || "$RUNTIME_IMAGE" == public.ecr.aws/* ]] + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" + metadata=$(docker image inspect "$RUNTIME_IMAGE") + resolved=$(jq -er '.[0].RepoDigests[0]' <<< "$metadata") + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + docker run --rm "$resolved" scripts/run_tracesbench_docker.py --help | grep -q -- '--data-root' \ + || { echo 'Comparison requires a runtime with --data-root support'; exit 1; } + fi + revision=$(jq -er '.[0].Config.Labels["org.opencontainers.image.revision"]' <<< "$metadata") + [[ "$revision" =~ ^[0-9a-f]{40}$ ]] + # Refuse runtimes without the trace runner and its validated dataset/report contract. + help=$(docker run --rm "${runtime_limits[@]}" "$resolved" scripts/run_tracesbench_docker.py --help) + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + grep -F -- '--baseline-run' <<< "$help" \ + || { echo 'Comparison requires a runtime with --baseline-run report support'; exit 1; } + fi + grep -F -- '--dataset {sanity,preflight,main}' <<< "$help" + grep -F -- '--runtime-cpus' <<< "$help" + grep -F -- '--runtime-memory' <<< "$help" + grep -F -- '--query-timeout-seconds' <<< "$help" + grep -F -- '--load-workers' <<< "$help" + grep -F -- '--targets' <<< "$help" + grep -F -- '--greptimedb-image' <<< "$help" + docker run --rm "${runtime_limits[@]}" "$resolved" scripts/check_run_tracesbench_docker.py + printf 'runtime_image=%s\no11ybench_ref=%s\nsystem_disk_gib=%s\n' "$resolved" "$revision" "$SYSTEM_DISK_GIB" >> "$GITHUB_OUTPUT" + provision: + needs: validate + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.ecs.outputs.label || steps.aws.outputs.label }} + instance_id: ${{ steps.ecs.outputs.instance_id || steps.aws.outputs.instance_id }} + runner_name: ${{ steps.ecs.outputs.runner_name || steps.aws.outputs.runner_name }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Create ephemeral ECS runner + if: inputs.provider == 'Aliyun' + id: ecs + uses: ./greptimedb/.github/actions/aliyun-ecs-create + with: + access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + access-key-secret: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} + vswitch-id: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} + security-group-id: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} + image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} + instance-type: ${{ needs.validate.outputs.instance_type }} + enable-docker: 'true' + system-disk-gib: ${{ needs.validate.outputs.system_disk_gib }} + ttl-hours: ${{ inputs.janitor_ttl_hours }} + resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} + runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + runner-gid: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Create ephemeral EC2 runner + if: inputs.provider == 'AWS' + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} + AWS_EC2_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py + benchmark: + needs: + - validate + - provision + runs-on: ${{ needs.provision.outputs.label }} + timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} + env: + TARGETS: ${{ needs.validate.outputs.targets }} + RESOLVED_RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} + RESOLVED_O11YBENCH_REF: ${{ needs.validate.outputs.o11ybench_ref }} + RUN_ROOT: ${{ inputs.benchmark_root }}/tracesbench-${{ github.run_id }}-${{ github.run_attempt }} + OWNER: o11ybench-traces-${{ github.run_id }}-${{ github.run_attempt }} + METADATA_ROOT: ${{ github.workspace }}/tracesbench-metadata + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark + - name: Checkout matching benchmark tools + uses: actions/checkout@v4 + with: + repository: GreptimeTeam/o11ybench + token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ref: ${{ env.RESOLVED_O11YBENCH_REF }} + path: o11ybench + persist-credentials: false + - name: Check runner and runtime identity + id: runtime + run: | + set -euo pipefail + test "$(git -C o11ybench rev-parse HEAD)" = "$RESOLVED_O11YBENCH_REF" + mkdir -p "$BENCHMARK_ROOT" "$METADATA_ROOT" + test "$(stat -f -c %T "$BENCHMARK_ROOT")" != tmpfs || { echo 'Benchmark root must be disk-backed'; exit 1; } + df -h "$BENCHMARK_ROOT" /var/lib/docker + minimum_gib=10 + if [[ "$DATASET" == main ]]; then minimum_gib=100; fi + for path in "$BENCHMARK_ROOT" /var/lib/docker; do + available=$(df -Pk "$path" | awk 'NR == 2 {print $4}') + ((available >= minimum_gib * 1024 * 1024)) || { echo "Need at least $minimum_gib GiB free on $path"; exit 1; } + done + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RESOLVED_RUNTIME_IMAGE" + docker image inspect "$RESOLVED_RUNTIME_IMAGE" > "$METADATA_ROOT/runtime-image.json" + test "$(jq -r '.[0].Config.Labels["org.opencontainers.image.revision"]' "$METADATA_ROOT/runtime-image.json")" = "$RESOLVED_O11YBENCH_REF" + - name: Pull database images before generating data + run: | + set -euo pipefail + IFS=, read -ra targets <<< "$TARGETS" + for target in "${targets[@]}"; do + if [[ "$target" == greptimedb ]]; then + image="greptime/greptimedb:$GREPTIMEDB_TAG" + else + image=$(jq -er '.target_identity.image | select(test("@sha256:[a-f0-9]{64}$"))' \ + "o11ybench/workloads/tracesbench/catalogs/$target.json") + fi + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$image" + docker image inspect "$image" > "$METADATA_ROOT/$target-image.json" + if [[ "$target" == greptimedb ]]; then + resolved=$(jq -er '.[0].RepoDigests[0] | select(test("@sha256:[a-f0-9]{64}$"))' "$METADATA_ROOT/$target-image.json") + printf 'GREPTIMEDB_IMAGE=%s\n' "$resolved" >> "$GITHUB_ENV" + fi + done + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + image="greptime/greptimedb:$BASELINE_GREPTIMEDB_TAG" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$image" + docker image inspect "$image" > "$METADATA_ROOT/greptimedb-baseline-image.json" + resolved=$(jq -er '.[0].RepoDigests[0] | select(test("@sha256:[a-f0-9]{64}$"))' "$METADATA_ROOT/greptimedb-baseline-image.json") + printf 'BASELINE_GREPTIMEDB_IMAGE=%s\n' "$resolved" >> "$GITHUB_ENV" + fi + - name: Generate selected dataset once + id: generate + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage generate \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Baseline GreptimeDB load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' && inputs.compare_greptimedb }} + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target greptimedb --port 14000 --greptimedb-image "$BASELINE_GREPTIMEDB_IMAGE" \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT/baseline" --data-root "$RUN_ROOT" --owner "$OWNER-base" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + python3 o11ybench/scripts/check_tracesbench_query_plans.py "$RUN_ROOT/baseline/runs/greptimedb" + - name: GreptimeDB load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' }} + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target greptimedb --port 14000 --greptimedb-image "$GREPTIMEDB_IMAGE" \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + python3 o11ybench/scripts/check_tracesbench_query_plans.py "$RUN_ROOT/runs/greptimedb" + - name: VictoriaTraces load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.victoriatraces == 'true' }} + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target victoriatraces --port 19428 \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Tempo load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.tempo == 'true' }} + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target tempo --port 13200 \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: GreptimeDB version comparison + if: always() && inputs.compare_greptimedb + run: | + python3 greptimedb/.github/scripts/benchmark-version-comparison.py report --kind traces --root "$RUN_ROOT" >> "$GITHUB_STEP_SUMMARY" + - name: Validate and render report + env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + run: | + comparison_args=() + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + comparison_args+=(--baseline-run "$RUN_ROOT/baseline/runs/greptimedb/run.json") + fi + python3 o11ybench/scripts/run_tracesbench_docker.py --stage report --targets "$TARGETS" "${comparison_args[@]}" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Clean up owned containers + if: always() && steps.runtime.outcome == 'success' + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage cleanup \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Clean up baseline containers + if: always() && steps.runtime.outcome == 'success' && inputs.compare_greptimedb + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage cleanup \ + --root "$RUN_ROOT/baseline" --owner "$OWNER-base" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Benchmark summary + if: always() + run: python3 greptimedb/.github/scripts/tracesbench-summary.py --root "$RUN_ROOT" --targets "$TARGETS" --section benchmark >> "$GITHUB_STEP_SUMMARY" + - name: Load / lifecycle summary + if: always() + run: python3 greptimedb/.github/scripts/tracesbench-summary.py --root "$RUN_ROOT" --targets "$TARGETS" --section lifecycle >> "$GITHUB_STEP_SUMMARY" + - name: Baseline load / lifecycle summary + if: always() && inputs.compare_greptimedb + run: | + echo '## Baseline GreptimeDB' >> "$GITHUB_STEP_SUMMARY" + python3 greptimedb/.github/scripts/tracesbench-summary.py \ + --root "$RUN_ROOT/baseline" --targets greptimedb --section lifecycle >> "$GITHUB_STEP_SUMMARY" + - name: Stage benchmark artifacts + if: always() + run: | + mkdir -p "$RUNNER_TEMP/tracesbench-artifacts" + if [[ -d "$METADATA_ROOT" ]]; then + rsync -a --no-owner --no-group "$METADATA_ROOT" "$RUNNER_TEMP/tracesbench-artifacts/" + fi + for entry in runs report report.json report.md version-comparison.json version-comparison.md; do + if [[ -e "$RUN_ROOT/$entry" ]]; then + rsync -a --no-owner --no-group "$RUN_ROOT/$entry" "$RUNNER_TEMP/tracesbench-artifacts/" + fi + done + if [[ -d "$RUN_ROOT/baseline/runs" ]]; then + mkdir -p "$RUNNER_TEMP/tracesbench-artifacts/baseline" + rsync -a --no-owner --no-group "$RUN_ROOT/baseline/runs" "$RUNNER_TEMP/tracesbench-artifacts/baseline/" + fi + - name: Upload HTML and results (no corpus) + if: always() + uses: actions/upload-artifact@v4 + with: + name: tracesbench-${{ inputs.dataset }}-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ runner.temp }}/tracesbench-artifacts + if-no-files-found: warn + retention-days: 14 + teardown: + needs: + - provision + - benchmark + if: ${{ always() && (needs.provision.outputs.instance_id != '' || (inputs.provider == 'AWS' && needs.provision.result != 'skipped')) }} + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Delete ECS and unregister runner + if: inputs.provider == 'Aliyun' + uses: ./greptimedb/.github/actions/aliyun-ecs-delete + with: + access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + access-key-secret: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} + instance-id: ${{ needs.provision.outputs.instance_id }} + runner-name: ${{ needs.provision.outputs.runner_name }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Delete EC2 and unregister runner + if: inputs.provider == 'AWS' + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py diff --git a/.github/workflows/vmbench-long-range.yml b/.github/workflows/vmbench-long-range.yml index 9d0088217ad..aaf2dd7f437 100644 --- a/.github/workflows/vmbench-long-range.yml +++ b/.github/workflows/vmbench-long-range.yml @@ -20,8 +20,18 @@ on: - 1mo default: 1w required: true + compare_greptimedb: + description: Compare baseline and candidate sequentially on the same host (adds one full GreptimeDB run) + type: boolean + default: false + required: false + baseline_greptimedb_tag: + description: Baseline Docker Hub greptime/greptimedb tag; required when comparison is enabled + type: string + default: v1.2.1 + required: false greptimedb_tag: - description: Docker Hub greptime/greptimedb tag + description: Candidate Docker Hub greptime/greptimedb tag type: string default: latest required: true @@ -31,14 +41,20 @@ on: default: v1.114.0@sha256:227ad53fd57a4b4430624898a1be0cac39431e1235fa0dc92f567cce38969344 required: true runtime_image: - description: Aliyun runtime image (tag or digest reference) + description: 'Runtime image: auto = AWS ECR Public or Aliyun ACR by provider; explicit tag/digest overrides' type: string - default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260918-545dc570@sha256:0ff03e7e690575e875bb85b3054a03486dd9359073546dc715da177153cf7a29 + default: auto required: true - ecs_instance_type: - description: ECS instance type (independent of Query Regression) + provider: + description: Cloud provider for the benchmark runner + type: choice + options: [Aliyun, AWS] + default: Aliyun + required: true + instance_type: + description: 'Instance type; auto = Aliyun: ecs.c9i.2xlarge, AWS: c7i.2xlarge (both 8 vCPU / 16 GiB). Enter a specific type to override.' type: string - default: ecs.c9i.2xlarge + default: auto required: true benchmark_root: description: Absolute disk-backed work directory (outside the o11ybench checkout) @@ -46,9 +62,9 @@ on: default: /home/runner/benchmark-data required: true system_disk_gib: - description: System disk GiB (20..2048) - type: number - default: 80 + description: 'System disk GiB: auto (or 0) = 80; enter 20..2048 to override' + type: string + default: auto required: true benchmark_timeout_minutes: description: Benchmark job timeout in minutes (1..360) @@ -73,12 +89,15 @@ on: permissions: contents: read env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + BASELINE_GREPTIMEDB_TAG: ${{ inputs.baseline_greptimedb_tag }} BENCHMARK_ROOT: ${{ inputs.benchmark_root }} RANGE_PRESET: ${{ inputs.range_preset }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} VICTORIAMETRICS_TAG: ${{ inputs.victoriametrics_tag }} - RUNTIME_IMAGE: ${{ inputs.runtime_image }} - ECS_INSTANCE_TYPE: ${{ inputs.ecs_instance_type }} + RUNTIME_IMAGE: ${{ inputs.runtime_image != 'auto' && inputs.runtime_image || inputs.provider == 'AWS' && 'public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime:b43785b56f42@sha256:e38811eb68caf3682ba56f5591a3e8e27e78e7983a923abb29fa2d7cd8fdbd46' || 'greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:b43785b56f42@sha256:e38811eb68caf3682ba56f5591a3e8e27e78e7983a923abb29fa2d7cd8fdbd46' }} + PROVIDER: ${{ inputs.provider }} + REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} @@ -89,6 +108,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 15 outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} system_disk_gib: ${{ steps.plan.outputs.system_disk_gib }} runtime_image: ${{ steps.plan.outputs.runtime_image }} o11ybench_ref: ${{ steps.plan.outputs.o11ybench_ref }} @@ -99,8 +119,19 @@ jobs: with: path: greptimedb persist-credentials: false - - name: Test existing ECS lifecycle helpers - run: python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py + - name: Validate version comparison + env: + RUN_GREPTIMEDB: ${{ inputs.run_greptimedb }} + run: | + python3 greptimedb/tests/perf/test_benchmark_version_comparison.py + python3 greptimedb/.github/scripts/benchmark-version-comparison.py validate + - name: Resolve CI runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" + - name: Test CI runner helpers + run: | + python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py + python3 greptimedb/tests/perf/test_ci_runner_scripts.py - name: Validate budgets and build warm plan id: plan env: @@ -108,11 +139,13 @@ jobs: RUN_VICTORIAMETRICS: ${{ inputs.run_victoriametrics }} run: | set -euo pipefail - [[ -n "$ECS_INSTANCE_TYPE" ]] [[ "$BENCHMARK_ROOT" == /* && "$BENCHMARK_ROOT" != / && "$BENCHMARK_ROOT" != */ ]] [[ "$BENCHMARK_ROOT" != /tmp && "$BENCHMARK_ROOT" != /tmp/* ]] [[ "$BENCHMARK_ROOT" != *$'\n'* && "$BENCHMARK_ROOT" != *$'\r'* ]] [[ "/${BENCHMARK_ROOT#/}/" != */../* && "/${BENCHMARK_ROOT#/}/" != */./* ]] + if [[ "$SYSTEM_DISK_GIB" == auto || "$SYSTEM_DISK_GIB" == 0 ]]; then + SYSTEM_DISK_GIB=80 + fi for value in "$SYSTEM_DISK_GIB" "$BENCHMARK_TIMEOUT_MINUTES" "$JANITOR_TTL_HOURS"; do [[ "$value" =~ ^[1-9][0-9]{0,3}$ ]] || { echo 'Resource budgets must be positive integers'; exit 1; } done @@ -120,8 +153,8 @@ jobs: ((BENCHMARK_TIMEOUT_MINUTES <= 360 && JANITOR_TTL_HOURS <= 168)) # Match the existing ECS workflow's provision, teardown and queue headroom. ((JANITOR_TTL_HOURS * 60 >= BENCHMARK_TIMEOUT_MINUTES + 120)) - [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* ]] - docker pull "$RUNTIME_IMAGE" + [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* || "$RUNTIME_IMAGE" == public.ecr.aws/* ]] + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" metadata=$(docker image inspect "$RUNTIME_IMAGE") resolved=$(jq -er '.[0].RepoDigests[0]' <<< "$metadata") revision=$(jq -er '.[0].Config.Labels["org.opencontainers.image.revision"]' <<< "$metadata") @@ -131,6 +164,9 @@ jobs: args=(--workload vmbench-long-range --range-preset "$RANGE_PRESET" --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --run-classification controlled_server_benchmark --no-open-browser) + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + args+=(--target "greptimedb:name=greptimedb-baseline,image=greptime/greptimedb:$BASELINE_GREPTIMEDB_TAG,url=http://127.0.0.1:14000/v1/prometheus/api/v1/query_range") + fi # Tags become target-spec fields; reject separators before passing them to the planner. for target in greptimedb victoriametrics; do case "$target" in @@ -159,15 +195,16 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 45 outputs: - label: ${{ steps.ecs.outputs.label }} - instance_id: ${{ steps.ecs.outputs.instance_id }} - runner_name: ${{ steps.ecs.outputs.runner_name }} + label: ${{ steps.ecs.outputs.label || steps.aws.outputs.label }} + instance_id: ${{ steps.ecs.outputs.instance_id || steps.aws.outputs.instance_id }} + runner_name: ${{ steps.ecs.outputs.runner_name || steps.aws.outputs.runner_name }} steps: - uses: actions/checkout@v4 with: path: greptimedb persist-credentials: false - name: Create ephemeral ECS runner + if: inputs.provider == 'Aliyun' id: ecs uses: ./greptimedb/.github/actions/aliyun-ecs-create with: @@ -178,13 +215,35 @@ jobs: vswitch-id: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} security-group-id: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} - instance-type: ${{ env.ECS_INSTANCE_TYPE }} + instance-type: ${{ needs.validate.outputs.instance_type }} enable-docker: 'true' system-disk-gib: ${{ needs.validate.outputs.system_disk_gib }} ttl-hours: ${{ inputs.janitor_ttl_hours }} resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} runner-gid: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Create ephemeral EC2 runner + if: inputs.provider == 'AWS' + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} + AWS_EC2_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py benchmark: needs: - validate @@ -202,6 +261,12 @@ jobs: TSDG_BATCH_SIZE: '100000' O11YBENCH_GENERATOR_CPUS: '4' steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark - name: Checkout matching benchmark tools uses: actions/checkout@v4 with: @@ -229,13 +294,13 @@ jobs: available=$(df -Pk "$path" | awk 'NR == 2 {print $4}') ((available >= 10 * 1024 * 1024)) || { echo "Need at least 10 GiB free on $path"; exit 1; } done - docker pull "$RESOLVED_RUNTIME_IMAGE" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RESOLVED_RUNTIME_IMAGE" docker image inspect "$RESOLVED_RUNTIME_IMAGE" > "$METADATA_ROOT/runtime-image.json" test "$(jq -r '.[0].Config.Labels["org.opencontainers.image.revision"]' "$METADATA_ROOT/runtime-image.json")" = "$RESOLVED_O11YBENCH_REF" plan="$METADATA_ROOT/warm.json" while IFS=$'\t' read -r target reference; do echo "Resolving $target: $reference" - docker pull "$reference" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$reference" docker image inspect "$reference" > "$METADATA_ROOT/$target-image.json" digest=$(jq -er '.[0].RepoDigests[0]' "$METADATA_ROOT/$target-image.json") jq --arg target "$target" --arg image "$digest" '.targets |= map(if .name == $target then .image = $image else . end)' "$plan" > "$plan.tmp" @@ -259,6 +324,10 @@ jobs: if: always() run: | docker ps -aq --filter "label=o11ybench.run-id=$RUN_ID-warm" | xargs -r docker rm -fv + - name: GreptimeDB version comparison + if: always() && inputs.compare_greptimedb + run: | + python3 greptimedb/.github/scripts/benchmark-version-comparison.py report --kind long-range --root "$ARTIFACT_ROOT-warm" >> "$GITHUB_STEP_SUMMARY" - name: Stage benchmark artifacts if: always() run: | @@ -283,7 +352,7 @@ jobs: needs: - provision - benchmark - if: ${{ always() && needs.provision.outputs.instance_id != '' }} + if: ${{ always() && (needs.provision.outputs.instance_id != '' || (inputs.provider == 'AWS' && needs.provision.result != 'skipped')) }} runs-on: ubuntu-latest timeout-minutes: 15 steps: @@ -292,6 +361,7 @@ jobs: path: greptimedb persist-credentials: false - name: Delete ECS and unregister runner + if: inputs.provider == 'Aliyun' uses: ./greptimedb/.github/actions/aliyun-ecs-delete with: access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} @@ -300,3 +370,20 @@ jobs: region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} instance-id: ${{ needs.provision.outputs.instance_id }} runner-name: ${{ needs.provision.outputs.runner_name }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Delete EC2 and unregister runner + if: inputs.provider == 'AWS' + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py diff --git a/.gitignore b/.gitignore index 90893f8aac1..dd60d96224e 100644 --- a/.gitignore +++ b/.gitignore @@ -52,8 +52,6 @@ venv/ tests-fuzz/artifacts/ tests-fuzz/corpus/ -# cargo-udeps reports -udeps-report.json # Nix .direnv diff --git a/AGENTS.md b/AGENTS.md index ef56c326805..20b2689a401 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,7 +25,7 @@ provides sub-second querying at PB scale with high cost efficiency. | Lint | `make clippy` (= `cargo clippy --workspace --all-targets --all-features -- -D warnings`) | | Type check | `make check` | -Toolchain: Rust nightly, Protobuf compiler (>= 3.15), C/C++ build essentials. +Toolchain: Rust (stable, pinned by `rust-toolchain.toml`), Protobuf compiler (>= 3.15), C/C++ build essentials. Install the test runner with `cargo install cargo-nextest --locked`. ## Repo map diff --git a/Cargo.lock b/Cargo.lock index 2c7f3333329..055bf712394 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -40,10 +40,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", - "cipher", + "cipher 0.4.4", "cpufeatures 0.2.17", ] +[[package]] +name = "aes" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f0f96ce78e38c3dc6d8948aa8163d06385be74000f3c7a95bf1eef35d3ea32" +dependencies = [ + "cipher 0.5.2", + "cpubits", + "cpufeatures 0.3.0", +] + [[package]] name = "aes-siv" version = "0.7.0" @@ -51,10 +62,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e08d0cdb774acd1e4dac11478b1a0c0d203134b2aab0ba25eb430de9b18f8b9" dependencies = [ "aead", - "aes", - "cipher", + "aes 0.8.4", + "cipher 0.4.4", "cmac", - "ctr", + "ctr 0.9.2", "dbl", "digest 0.10.7", "zeroize", @@ -219,7 +230,7 @@ checksum = "d301b3b94cb4b2f23d7917810addbbaff90738e0ca2be692bd027e70d7e0330c" [[package]] name = "api" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow-schema 59.2.0", "common-base", @@ -948,7 +959,7 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "auth" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -1400,6 +1411,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-padding" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "710f1dd022ef4e93f8a438b4ba958de7f64308434fa6a87104481645cc30068b" +dependencies = [ + "hybrid-array", +] + [[package]] name = "blocking" version = "1.6.1" @@ -1534,9 +1554,9 @@ checksum = "40e38929add23cdf8a366df9b0e088953150724bcbe5fc330b0d8eb3b328eec8" [[package]] name = "bumpalo" -version = "3.19.0" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "bytecheck" @@ -1598,7 +1618,7 @@ dependencies = [ [[package]] name = "cache" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "catalog", "common-error", @@ -1664,7 +1684,7 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "catalog" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow 59.2.0", @@ -1673,7 +1693,6 @@ dependencies = [ "async-trait", "bytes", "cache", - "catalog", "chrono", "common-base", "common-catalog", @@ -1701,7 +1720,6 @@ dependencies = [ "lazy_static", "meta-client", "moka", - "object-store", "partition", "paste", "prometheus 0.14.0", @@ -1732,7 +1750,16 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6" dependencies = [ - "cipher", + "cipher 0.4.4", +] + +[[package]] +name = "cbc" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce2dc9ee5f88d11e0beb842c88b33c8a5cf0d1329c4b19494af42b07dbfe8896" +dependencies = [ + "cipher 0.5.2", ] [[package]] @@ -1747,15 +1774,6 @@ dependencies = [ "shlex 2.0.1", ] -[[package]] -name = "cedarwood" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0524a528a6a0288df1863c3c20fe92c301875b4941e7b6c4b394ab08c5a4c55" -dependencies = [ - "smallvec", -] - [[package]] name = "census" version = "0.4.2" @@ -1777,7 +1795,7 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "738b8d467867f80a71351933f70461f5b56f24d5c93e0cf216e59229c968d330" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -1829,7 +1847,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" dependencies = [ "cfg-if", - "cipher", + "cipher 0.4.4", "cpufeatures 0.2.17", ] @@ -1852,7 +1870,7 @@ checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" dependencies = [ "aead", "chacha20 0.9.1", - "cipher", + "cipher 0.4.4", "poly1305", "zeroize", ] @@ -1900,7 +1918,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7d8d1efd5109b9c1cd3b7966bd071cdfb53bb6eb0b22a473a68c2f70a11a1eb" dependencies = [ "parse-zoneinfo", - "phf_codegen 0.12.1", + "phf_codegen", "phf_shared 0.12.1", "uncased", ] @@ -1945,10 +1963,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" dependencies = [ "crypto-common 0.1.6", - "inout", + "inout 0.1.4", "zeroize", ] +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "block-buffer 0.12.0", + "crypto-common 0.2.2", + "inout 0.2.2", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -1957,7 +1986,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" dependencies = [ "glob", "libc", - "libloading", + "libloading 0.8.8", ] [[package]] @@ -2002,14 +2031,12 @@ checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" [[package]] name = "cli" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-stream", "async-trait", - "auth", "base64 0.22.1", "cache", - "catalog", "chrono", "clap", "client", @@ -2022,13 +2049,10 @@ dependencies = [ "common-macro", "common-meta", "common-procedure", - "common-query", - "common-recordbatch", "common-runtime", "common-telemetry", "common-test-util", "common-time", - "common-version", "common-wal", "datatypes", "etcd-client", @@ -2048,7 +2072,6 @@ dependencies = [ "serde", "serde_json", "servers", - "session", "snafu 0.8.6", "sql", "standalone", @@ -2056,14 +2079,14 @@ dependencies = [ "table", "tempfile", "tokio", - "tracing-appender", + "tokio-util", "url", "uuid", ] [[package]] name = "client" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arc-swap", @@ -2074,7 +2097,6 @@ dependencies = [ "common-catalog", "common-error", "common-grpc", - "common-grpc-expr", "common-macro", "common-meta", "common-query", @@ -2096,11 +2118,10 @@ dependencies = [ "serde_json", "snafu 0.8.6", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.4.0-alpha.0", "tokio", "tokio-stream", "tonic 0.14.2", - "tracing", ] [[package]] @@ -2129,7 +2150,7 @@ version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8543454e3c3f5126effff9cd44d562af4e31fb8ce1cc0d3dcd8f084515dbc1aa" dependencies = [ - "cipher", + "cipher 0.4.4", "dbl", "digest 0.10.7", ] @@ -2145,15 +2166,13 @@ dependencies = [ [[package]] name = "cmd" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", "auth", - "base64 0.22.1", "cache", "catalog", - "chrono", "clap", "cli", "client", @@ -2170,7 +2189,6 @@ dependencies = [ "common-options", "common-procedure", "common-query", - "common-recordbatch", "common-runtime", "common-stat", "common-telemetry", @@ -2184,7 +2202,6 @@ dependencies = [ "datafusion-physical-plan", "datanode", "datatypes", - "either", "etcd-client", "file-engine", "flow", @@ -2192,7 +2209,6 @@ dependencies = [ "futures", "hex", "human-panic", - "humantime", "lazy_static", "log-store", "meta-client", @@ -2216,20 +2232,17 @@ dependencies = [ "serde", "serde_json", "servers", - "session", "similar-asserts", "smallvec", "snafu 0.8.6", "sqlparser", "standalone", "store-api", - "table", "temp-env", "tempfile", "tikv-jemallocator", "tokio", "toml 0.8.23", - "tonic 0.14.2", "tracing-appender", ] @@ -2247,18 +2260,7 @@ checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", -] - -[[package]] -name = "codespan-reporting" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" -dependencies = [ - "serde", - "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -2300,7 +2302,7 @@ dependencies = [ [[package]] name = "common-base" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "anymap2", @@ -2321,7 +2323,7 @@ dependencies = [ [[package]] name = "common-batcher" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "dashmap", "futures", @@ -2330,14 +2332,14 @@ dependencies = [ [[package]] name = "common-catalog" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "const_format", ] [[package]] name = "common-config" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-error", @@ -2355,13 +2357,12 @@ dependencies = [ "serde_json", "snafu 0.8.6", "temp-env", - "tempfile", "toml 0.8.23", ] [[package]] name = "common-datasource" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -2388,6 +2389,7 @@ dependencies = [ "paste", "regex", "serde", + "serde_json", "snafu 0.8.6", "strum 0.27.1", "tokio", @@ -2397,7 +2399,7 @@ dependencies = [ [[package]] name = "common-decimal" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "bigdecimal 0.4.8", "common-error", @@ -2410,7 +2412,7 @@ dependencies = [ [[package]] name = "common-error" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-macro", "http 1.5.0", @@ -2422,7 +2424,7 @@ dependencies = [ [[package]] name = "common-event-recorder" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -2447,7 +2449,7 @@ dependencies = [ [[package]] name = "common-frontend" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -2456,40 +2458,35 @@ dependencies = [ "common-grpc", "common-macro", "common-meta", - "common-time", "greptime-proto", - "humantime", + "lazy_static", "meta-client", + "prometheus 0.14.0", "serde", "serde_json", "snafu 0.8.6", - "tokio", "tonic 0.14.2", ] [[package]] name = "common-function" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", "approx 0.5.1", - "arc-swap", "arrow 59.2.0", - "arrow-cast 59.2.0", "arrow-schema 59.2.0", "async-trait", "axum 0.8.4", "bincode", "catalog", - "chrono", "common-base", "common-catalog", "common-error", "common-macro", "common-meta", "common-query", - "common-runtime", "common-telemetry", "common-time", "common-version", @@ -2514,7 +2511,6 @@ dependencies = [ "memchr", "mito-codec", "nalgebra", - "num", "num-traits", "paste", "pretty_assertions", @@ -2536,7 +2532,7 @@ dependencies = [ [[package]] name = "common-greptimedb-telemetry" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "common-runtime", @@ -2553,7 +2549,7 @@ dependencies = [ [[package]] name = "common-grpc" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-flight", @@ -2588,10 +2584,9 @@ dependencies = [ [[package]] name = "common-grpc-expr" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", - "common-base", "common-catalog", "common-error", "common-macro", @@ -2599,7 +2594,6 @@ dependencies = [ "common-sql", "common-time", "datatypes", - "paste", "prost 0.14.1", "snafu 0.8.6", "store-api", @@ -2608,7 +2602,7 @@ dependencies = [ [[package]] name = "common-macro" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "greptime-proto", "once_cell", @@ -2619,7 +2613,7 @@ dependencies = [ [[package]] name = "common-mem-prof" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "anyhow", "common-error", @@ -2635,7 +2629,7 @@ dependencies = [ [[package]] name = "common-memory-manager" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-error", "common-macro", @@ -2647,7 +2641,7 @@ dependencies = [ [[package]] name = "common-meta" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "anymap2", "api", @@ -2675,7 +2669,6 @@ dependencies = [ "common-wal", "common-workload", "datafusion-common", - "datafusion-expr", "datatypes", "deadpool", "deadpool-postgres", @@ -2684,7 +2677,6 @@ dependencies = [ "flexbuffers", "futures", "futures-util", - "greptime-proto", "hex", "humantime-serde", "itertools 0.14.0", @@ -2720,7 +2712,7 @@ dependencies = [ [[package]] name = "common-options" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-grpc", @@ -2731,11 +2723,11 @@ dependencies = [ [[package]] name = "common-plugins" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" [[package]] name = "common-pprof" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-error", "common-macro", @@ -2746,7 +2738,7 @@ dependencies = [ [[package]] name = "common-procedure" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-stream", @@ -2775,17 +2767,16 @@ dependencies = [ [[package]] name = "common-procedure-test" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "common-procedure", - "snafu 0.8.6", "tokio", ] [[package]] name = "common-query" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -2815,7 +2806,7 @@ dependencies = [ [[package]] name = "common-recordbatch" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "arc-swap", "common-base", @@ -2840,7 +2831,7 @@ dependencies = [ [[package]] name = "common-runtime" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "catio", @@ -2852,13 +2843,11 @@ dependencies = [ "lazy_static", "num_cpus", "once_cell", - "parking_lot 0.12.4", "paste", "pin-project", "prometheus 0.14.0", "ratelimit", "serde", - "serde_json", "snafu 0.8.6", "tempfile", "tokio", @@ -2870,7 +2859,7 @@ dependencies = [ [[package]] name = "common-session" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "serde", "strum 0.27.1", @@ -2878,14 +2867,13 @@ dependencies = [ [[package]] name = "common-sql" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow-schema 59.2.0", "common-base", "common-decimal", "common-error", "common-macro", - "common-telemetry", "common-time", "datafusion-sql", "datatypes", @@ -2898,12 +2886,11 @@ dependencies = [ [[package]] name = "common-stat" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-runtime", "common-telemetry", - "lazy_static", "nix 0.30.1", "num_cpus", "prometheus 0.14.0", @@ -2913,7 +2900,7 @@ dependencies = [ [[package]] name = "common-telemetry" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "backtrace", "common-base", @@ -2943,7 +2930,7 @@ dependencies = [ [[package]] name = "common-test-util" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-grpc", "common-query", @@ -2955,7 +2942,7 @@ dependencies = [ [[package]] name = "common-time" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow 59.2.0", "chrono", @@ -2973,7 +2960,7 @@ dependencies = [ [[package]] name = "common-version" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "cargo-manifest", "const_format", @@ -2983,7 +2970,7 @@ dependencies = [ [[package]] name = "common-wal" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-error", @@ -3006,7 +2993,7 @@ dependencies = [ [[package]] name = "common-workload" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "common-telemetry", "serde", @@ -3266,6 +3253,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "cpubits" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -3286,9 +3279,9 @@ dependencies = [ [[package]] name = "crc" -version = "3.3.0" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9710d3b3739c2e349eb44fe848ad0b7c8cb1e42bd87ee49371df2f7acaf3e675" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" dependencies = [ "crc-catalog", ] @@ -3486,7 +3479,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9d6cf87adf719ddf43a805e92c6870a531aedda35ff640442cbaf8674e141e1" dependencies = [ "aead", - "cipher", + "cipher 0.4.4", "generic-array", "poly1305", "salsa20", @@ -3531,7 +3524,16 @@ version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" dependencies = [ - "cipher", + "cipher 0.4.4", +] + +[[package]] +name = "ctr" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baaca1c4b237092596f64d571e9db6ce4109c4ef9742e27590f1709594461f21" +dependencies = [ + "cipher 0.5.2", ] [[package]] @@ -3543,68 +3545,6 @@ dependencies = [ "cmov", ] -[[package]] -name = "cxx" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7620f6cfc4dcca21f2b085b7a890e16c60fd66f560cd69ee60594908dc72ab1" -dependencies = [ - "cc", - "cxx-build", - "cxxbridge-cmd", - "cxxbridge-flags", - "cxxbridge-macro", - "foldhash 0.2.0", - "link-cplusplus", -] - -[[package]] -name = "cxx-build" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a9bc1a22964ff6a355fbec24cf68266a0ed28f8b84c0864c386474ea3d0e479" -dependencies = [ - "cc", - "codespan-reporting 0.13.1", - "indexmap 2.14.2", - "proc-macro2", - "quote", - "scratch", - "syn 2.0.117", -] - -[[package]] -name = "cxxbridge-cmd" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f29a879d35f7906e3c9b77d7a1005a6a0787d330c09dfe4ffb5f617728cb44" -dependencies = [ - "clap", - "codespan-reporting 0.13.1", - "indexmap 2.14.2", - "proc-macro2", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "cxxbridge-flags" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d67109015f93f683e364085aa6489a5b2118b4a40058482101d699936a7836d6" - -[[package]] -name = "cxxbridge-macro" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d187e019e7b05a1f3e69a8396b70800ee867aa9fc2ab972761173ccee03742df" -dependencies = [ - "indexmap 2.14.2", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "daemonize" version = "0.5.0" @@ -3782,7 +3722,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3834,7 +3774,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3858,7 +3798,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3881,7 +3821,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3907,7 +3847,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "futures", "log", @@ -3917,7 +3857,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-compression", @@ -3953,7 +3893,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3977,7 +3917,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4000,7 +3940,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4023,7 +3963,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4055,12 +3995,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" [[package]] name = "datafusion-execution" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4085,7 +4025,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4109,7 +4049,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4120,7 +4060,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4151,7 +4091,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4171,7 +4111,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4182,7 +4122,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-ord 59.2.0", @@ -4206,7 +4146,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4221,7 +4161,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4237,7 +4177,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4246,7 +4186,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "datafusion-doc", "quote", @@ -4256,7 +4196,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "chrono", @@ -4304,7 +4244,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4327,7 +4267,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4341,7 +4281,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "chrono", @@ -4358,7 +4298,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4377,7 +4317,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-data 59.2.0", @@ -4414,7 +4354,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-catalog", @@ -4440,7 +4380,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4450,7 +4390,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -4460,7 +4400,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4475,7 +4415,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow-schema 59.2.0", "async-trait", @@ -4489,7 +4429,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "bigdecimal 0.4.8", @@ -4508,7 +4448,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "async-recursion", "async-trait", @@ -4526,13 +4466,14 @@ dependencies = [ [[package]] name = "datanode" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-flight", "async-trait", "bytes", "cache", + "catalog", "client", "common-base", "common-catalog", @@ -4601,7 +4542,7 @@ checksum = "c286de4e81ea2590afc24d754e0f83810c566f50a1388fa75ebd57928c0d9745" [[package]] name = "datatypes" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow 59.2.0", "arrow-array 59.2.0", @@ -4854,6 +4795,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "des" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "916a94e407b54f9034d71dd748234cd1e516ced6284009906ae246f177eafe5a" +dependencies = [ + "cipher 0.5.2", +] + [[package]] name = "diff" version = "0.1.13" @@ -5324,7 +5274,7 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "file-engine" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -5332,7 +5282,6 @@ dependencies = [ "common-datasource", "common-error", "common-macro", - "common-query", "common-recordbatch", "common-telemetry", "common-test-util", @@ -5466,25 +5415,20 @@ checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" [[package]] name = "flow" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow 59.2.0", "arrow-flight", "arrow-schema 59.2.0", - "async-recursion", "async-trait", - "bytes", "catalog", "chrono", "client", "common-base", "common-catalog", "common-config", - "common-datasource", - "common-decimal", "common-error", - "common-frontend", "common-function", "common-grpc", "common-macro", @@ -5497,28 +5441,20 @@ dependencies = [ "common-telemetry", "common-time", "common-version", - "config", "datafusion", "datafusion-common", "datafusion-expr", "datafusion-physical-expr", - "datafusion-substrait", "datatypes", - "enum_dispatch", "futures", "get-size2", "greptime-proto", "hostname 0.4.1", - "http 1.5.0", "humantime-serde", "itertools 0.14.0", "lazy_static", "meta-client", - "nom 7.1.3", - "num-traits", - "num_cpus", "operator", - "partition", "pretty_assertions", "prometheus 0.14.0", "prost 0.14.1", @@ -5528,12 +5464,10 @@ dependencies = [ "serde_json", "servers", "session", - "smallvec", "snafu 0.8.6", "sql", "store-api", - "strum 0.27.1", - "substrait 1.3.0-alpha.1", + "substrait 1.4.0-alpha.0", "table", "tokio", "tokio-stream", @@ -5630,10 +5564,9 @@ dependencies = [ [[package]] name = "frontend" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", - "arc-swap", "arrow-flight", "async-stream", "async-trait", @@ -5657,7 +5590,6 @@ dependencies = [ "common-memory-manager", "common-meta", "common-options", - "common-procedure", "common-query", "common-recordbatch", "common-runtime", @@ -5672,13 +5604,12 @@ dependencies = [ "datatypes", "futures", "hostname 0.4.1", - "humantime", "humantime-serde", "hyper-util", "lazy_static", "log-query", "meta-client", - "num_cpus", + "moka", "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", "operator", @@ -5689,7 +5620,6 @@ dependencies = [ "promql-parser", "prost 0.14.1", "query", - "rand 0.9.4", "reqwest 0.13.4", "serde", "serde_json", @@ -5708,7 +5638,6 @@ dependencies = [ "tonic 0.14.2", "tower 0.5.2", "tracing", - "uuid", ] [[package]] @@ -5953,6 +5882,34 @@ dependencies = [ "byteorder", ] +[[package]] +name = "g2gen" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5a7e0eb46f83a20260b850117d204366674e85d3a908d90865c78df9a6b1dfc" +dependencies = [ + "g2poly", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "g2p" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "539e2644c030d3bf4cd208cb842d2ce2f80e82e6e8472390bcef83ceba0d80ad" +dependencies = [ + "g2gen", + "g2poly", +] + +[[package]] +name = "g2poly" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "312d2295c7302019c395cfb90dacd00a82a2eabd700429bba9c7a3f38dbbe11b" + [[package]] name = "generator" version = "0.8.5" @@ -6130,7 +6087,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=0c4beffaa2b3968034915867593a717c6ab0c6b4#0c4beffaa2b3968034915867593a717c6ab0c6b4" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=9533078d6fd4bdaddbca0750cb9bf7ad64b2e568#9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -6301,6 +6258,47 @@ dependencies = [ "hashbrown 0.15.4", ] +[[package]] +name = "hdfs-native" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd181084003308224efddf737417832186839ce2882d2468ddad0114cfb3551" +dependencies = [ + "aes 0.9.3", + "base64 0.22.1", + "bitflags 2.12.1", + "bumpalo", + "bytes", + "cbc 0.2.1", + "chrono", + "cipher 0.5.2", + "crc", + "ctr 0.10.1", + "des", + "dns-lookup", + "futures", + "g2p", + "hex", + "hmac 0.13.0", + "libc", + "libloading 0.9.0", + "log", + "md-5 0.11.0", + "num-traits", + "once_cell", + "prost 0.14.1", + "prost-types 0.14.1", + "rand 0.10.1", + "regex", + "roxmltree", + "socket2 0.6.4", + "thiserror 2.0.17", + "tokio", + "url", + "uuid", + "whoami 2.1.3", +] + [[package]] name = "hdrhistogram" version = "7.5.4" @@ -6683,7 +6681,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -6754,7 +6752,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.57.0", + "windows-core 0.61.2", ] [[package]] @@ -6945,8 +6943,9 @@ dependencies = [ [[package]] name = "index" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ + "ahash 0.8.12", "async-trait", "asynchronous-codec", "bytemuck", @@ -6958,7 +6957,6 @@ dependencies = [ "common-telemetry", "common-test-util", "criterion 0.7.0", - "datatypes", "fastbloom", "fst", "futures", @@ -6967,7 +6965,6 @@ dependencies = [ "jieba-rs", "lazy_static", "mockall", - "nalgebra", "pin-project", "prost 0.14.1", "puffin", @@ -6985,7 +6982,6 @@ dependencies = [ "tempfile", "tokio", "tokio-util", - "usearch", "uuid", ] @@ -7132,10 +7128,20 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" dependencies = [ - "block-padding", + "block-padding 0.3.3", "generic-array", ] +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "block-padding 0.4.2", + "hybrid-array", +] + [[package]] name = "instant" version = "0.1.13" @@ -7197,7 +7203,7 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96e4f67dbfc0f75d7b65953ecf0be3fd84ee0cb1ae72a00a4aa9a2f5518a2c80" dependencies = [ - "aes", + "aes 0.8.4", ] [[package]] @@ -7334,25 +7340,19 @@ dependencies = [ [[package]] name = "jieba-macros" -version = "0.10.1" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46adade69b634535a8f495cf87710ed893cff53e1dbc9dd750c2ab81c5defb82" -dependencies = [ - "phf_codegen 0.13.1", -] +checksum = "2987cf45b1bfdb4ff1f6a814b1aa2ad6361f4ff43bf353a9bd29f391a9568f70" [[package]] name = "jieba-rs" -version = "0.10.1" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11b53580aaa8ec8b713da271da434f8947409242c537a9ab3f7b76bdbb19e8a9" +checksum = "0a3a7d02ea31e6fc6f2fe10495dc7c5efcf381e1f21d7172f58ba792efa1c2fd" dependencies = [ "bytecount", - "cedarwood", "include-flate", "jieba-macros", - "phf 0.13.1", - "regex", "rustc-hash", ] @@ -7946,6 +7946,16 @@ dependencies = [ "windows-targets 0.53.5", ] +[[package]] +name = "libloading" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +dependencies = [ + "cfg-if", + "windows-link 0.2.1", +] + [[package]] name = "liblzma" version = "0.4.6" @@ -8006,15 +8016,6 @@ dependencies = [ "vcpkg", ] -[[package]] -name = "link-cplusplus" -version = "1.0.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f78c730aaa7d0b9336a299029ea49f9ee53b0ed06e9202e8cb7db9bae7b8c82" -dependencies = [ - "cc", -] - [[package]] name = "link-section" version = "0.18.1" @@ -8087,7 +8088,7 @@ checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "log-query" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "chrono", "common-error", @@ -8099,7 +8100,7 @@ dependencies = [ [[package]] name = "log-store" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-stream", "async-trait", @@ -8121,7 +8122,6 @@ dependencies = [ "derive_builder 0.20.2", "futures", "futures-util", - "itertools 0.14.0", "lazy_static", "object-store", "pin-project", @@ -8452,7 +8452,7 @@ dependencies = [ [[package]] name = "meta-client" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -8465,34 +8465,28 @@ dependencies = [ "common-options", "common-telemetry", "common-time", - "datatypes", "futures", - "futures-util", "humantime-serde", "meta-srv", "rand 0.9.4", "serde", "serde_json", - "session", "snafu 0.8.6", "tokio", "tokio-stream", "tonic 0.14.2", - "tower 0.5.2", "tracing", "tracing-subscriber", ] [[package]] name = "meta-srv" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", "axum 0.8.4", - "axum-extra", "axum-macros", - "bytes", "chrono", "clap", "client", @@ -8525,19 +8519,14 @@ dependencies = [ "futures", "h2 0.3.26", "hostname 0.4.1", - "http-body-util", - "humantime", "humantime-serde", "hyper-util", "itertools 0.14.0", "lazy_static", - "local-ip-address", "once_cell", "ordered-float 4.6.0", - "parking_lot 0.12.4", "partition", "prometheus 0.14.0", - "prost 0.14.1", "rand 0.9.4", "regex", "rskafka", @@ -8555,7 +8544,6 @@ dependencies = [ "toml 0.8.23", "tonic 0.14.2", "tower 0.5.2", - "tower-http 0.6.11", "tracing", "tracing-subscriber", "typetag", @@ -8584,7 +8572,7 @@ dependencies = [ [[package]] name = "metric-engine" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "aquamarine", @@ -8603,7 +8591,6 @@ dependencies = [ "common-recordbatch", "common-runtime", "common-telemetry", - "common-test-util", "common-time", "common-wal", "criterion 0.7.0", @@ -8685,7 +8672,7 @@ dependencies = [ [[package]] name = "mito-codec" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "bytes", @@ -8710,7 +8697,7 @@ dependencies = [ [[package]] name = "mito2" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "aquamarine", @@ -8723,9 +8710,7 @@ dependencies = [ "bytes", "chrono", "common-base", - "common-config", "common-datasource", - "common-decimal", "common-error", "common-function", "common-grpc", @@ -8764,7 +8749,6 @@ dependencies = [ "object-store", "parquet", "partition", - "paste", "pin-project", "prometheus 0.14.0", "prost 0.14.1", @@ -8772,15 +8756,12 @@ dependencies = [ "rand 0.9.4", "rayon", "regex", - "roaring", "rskafka", "rstest", "rstest_reuse", - "scopeguard", "serde", "serde_json", "serde_with", - "session", "smallvec", "snafu 0.8.6", "store-api", @@ -9465,7 +9446,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d" dependencies = [ - "proc-macro-crate 1.3.1", + "proc-macro-crate 3.3.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9559,13 +9540,11 @@ dependencies = [ [[package]] name = "object-store" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "anyhow", - "async-trait", "bytes", "cap-std", - "chrono", "common-base", "common-error", "common-macro", @@ -9574,14 +9553,14 @@ dependencies = [ "common-test-util", "derive_builder 0.20.2", "futures", + "hdfs-native", "humantime-serde", - "lazy_static", + "object-store", "object_store", "object_store_opendal", "opendal", "opendal-http-transport-reqwest", "prometheus 0.14.0", - "rand 0.9.4", "reqwest 0.13.4", "serde", "snafu 0.8.6", @@ -9650,7 +9629,7 @@ version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2cc40678e045ff4eb1666ea6c0f994b133c31f673c09aed292261b6d5b6963a0" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -9723,6 +9702,7 @@ dependencies = [ "opendal-service-azblob", "opendal-service-fs", "opendal-service-gcs", + "opendal-service-hdfs-native", "opendal-service-http", "opendal-service-mysql", "opendal-service-oss", @@ -9913,6 +9893,20 @@ dependencies = [ "uuid", ] +[[package]] +name = "opendal-service-hdfs-native" +version = "0.59.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b84340e11a3bcd7f1f0b55df032d040b2907483c546c3021485aba749c3072c" +dependencies = [ + "bytes", + "futures", + "hdfs-native", + "log", + "opendal-core", + "serde", +] + [[package]] name = "opendal-service-http" version = "0.59.2" @@ -9978,7 +9972,7 @@ dependencies = [ [[package]] name = "opensrv-mysql" version = "0.8.0" -source = "git+https://github.com/GreptimeTeam/opensrv?rev=f3135e5f00401e06075cf241c0cc26489e45b20a#f3135e5f00401e06075cf241c0cc26489e45b20a" +source = "git+https://github.com/GreptimeTeam/opensrv?rev=1f3d88ccaafb5721be6cfcc3ccbca48642426119#1f3d88ccaafb5721be6cfcc3ccbca48642426119" dependencies = [ "async-trait", "byteorder", @@ -10091,12 +10085,11 @@ dependencies = [ [[package]] name = "operator" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", "arrow 59.2.0", - "arrow-ipc 59.2.0", "async-stream", "async-trait", "axum 0.8.4", @@ -10143,7 +10136,6 @@ dependencies = [ "path-slash", "percent-encoding", "prometheus 0.14.0", - "prost 0.14.1", "prost-types 0.14.1", "query", "regex", @@ -10155,7 +10147,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.4.0-alpha.0", "table", "tokio", "tokio-util", @@ -10489,22 +10481,18 @@ checksum = "e3c406c9e2aa74554e662d2c2ee11cd3e73756988800be7e6f5eddb16fed4699" [[package]] name = "partition" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", - "async-trait", "common-base", "common-error", "common-macro", "common-meta", - "common-query", - "common-telemetry", "criterion 0.7.0", "datafusion-common", "datafusion-expr", "datafusion-physical-expr", "datatypes", - "itertools 0.14.0", "moka", "rand 0.8.5", "serde", @@ -10777,20 +10765,10 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "efbdcb6f01d193b17f0b9c3360fa7e0e620991b193ff08702f78b3ce365d7e61" dependencies = [ - "phf_generator 0.12.1", + "phf_generator", "phf_shared 0.12.1", ] -[[package]] -name = "phf_codegen" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49aa7f9d80421bca176ca8dbfebe668cc7a2684708594ec9f3c0db0805d5d6e1" -dependencies = [ - "phf_generator 0.13.1", - "phf_shared 0.13.1", -] - [[package]] name = "phf_generator" version = "0.12.1" @@ -10801,23 +10779,13 @@ dependencies = [ "phf_shared 0.12.1", ] -[[package]] -name = "phf_generator" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "135ace3a761e564ec88c03a77317a7c6b80bb7f7135ef2544dbe054243b89737" -dependencies = [ - "fastrand", - "phf_shared 0.13.1", -] - [[package]] name = "phf_macros" version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d713258393a82f091ead52047ca779d37e5766226d009de21696c4e667044368" dependencies = [ - "phf_generator 0.12.1", + "phf_generator", "phf_shared 0.12.1", "proc-macro2", "quote", @@ -10887,35 +10855,29 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pipeline" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", "arrow 59.2.0", "arrow-schema 59.2.0", - "async-trait", "catalog", "chrono", "chrono-tz", "common-catalog", "common-error", - "common-function", + "common-frontend", "common-macro", "common-meta", "common-query", "common-recordbatch", - "common-runtime", "common-telemetry", "common-time", "criterion 0.7.0", - "crossbeam-utils", "csv", - "dashmap", "datafusion", "datafusion-common", "datafusion-expr", - "datafusion-functions", - "datafusion-physical-expr", "datatypes", "dyn-fmt", "enum_dispatch", @@ -10930,10 +10892,8 @@ dependencies = [ "once_cell", "operator", "ordered-float 4.6.0", - "paste", "prometheus 0.14.0", "query", - "rayon", "regex", "serde", "serde_json", @@ -10975,8 +10935,8 @@ version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" dependencies = [ - "aes", - "cbc", + "aes 0.8.4", + "cbc 0.1.2", "der", "pbkdf2", "scrypt", @@ -11051,7 +11011,7 @@ dependencies = [ [[package]] name = "plugins" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "auth", "catalog", @@ -11394,7 +11354,7 @@ dependencies = [ [[package]] name = "promql" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "async-trait", @@ -11496,7 +11456,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11516,8 +11476,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", - "itertools 0.10.5", + "heck 0.5.0", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -11565,7 +11525,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11578,7 +11538,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11811,7 +11771,7 @@ dependencies = [ [[package]] name = "puffin" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-compression", "async-trait", @@ -11863,11 +11823,10 @@ dependencies = [ [[package]] name = "query" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", - "arc-swap", "arrow 59.2.0", "arrow-schema 59.2.0", "async-recursion", @@ -11901,12 +11860,10 @@ dependencies = [ "datafusion-proto", "datafusion-sql", "datatypes", - "either", "fastrand", "futures", "futures-util", "greptime-proto", - "humantime", "itertools 0.14.0", "jsonb", "lazy_static", @@ -11915,11 +11872,8 @@ dependencies = [ "meter-macros", "moka", "nalgebra", - "num", - "num-traits", "object-store", "once_cell", - "parking_lot 0.12.4", "partition", "paste", "pretty_assertions", @@ -11936,10 +11890,9 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.4.0-alpha.0", "table", "tokio", - "tokio-stream", "toml 0.8.23", "tracing", "uuid", @@ -13178,7 +13131,7 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -13265,12 +13218,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" -[[package]] -name = "scratch" -version = "1.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d68f2ec51b097e4c1a75b681a8bec621909b5e91f15bb7b840c4f2f7b01148b2" - [[package]] name = "scrypt" version = "0.11.0" @@ -13558,7 +13505,7 @@ dependencies = [ [[package]] name = "servers" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -13602,7 +13549,6 @@ dependencies = [ "common-version", "criterion 0.5.1", "csv", - "dashmap", "datafusion", "datafusion-common", "datafusion-expr", @@ -13616,6 +13562,7 @@ dependencies = [ "hex", "hostname 0.3.1", "http 1.5.0", + "http-body 1.0.1", "humantime", "humantime-serde", "hyper 1.6.0", @@ -13702,7 +13649,7 @@ dependencies = [ [[package]] name = "session" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -13719,6 +13666,7 @@ dependencies = [ "datafusion-common", "derive_builder 0.20.2", "derive_more", + "prost 0.14.1", "snafu 0.8.6", "sql", "uuid", @@ -13976,7 +13924,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -14062,33 +14010,28 @@ dependencies = [ [[package]] name = "sql" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-buffer 59.2.0", + "arrow-schema 59.2.0", "chrono", "common-base", "common-catalog", "common-datasource", - "common-decimal", "common-error", "common-macro", "common-query", "common-sql", - "common-telemetry", "common-time", "datafusion", "datafusion-common", "datafusion-expr", - "datafusion-physical-expr", "datafusion-sql", "datatypes", "either", - "hex", - "humantime", "iso8601", "itertools 0.14.0", - "jsonb", "lazy_static", "pretty_assertions", "promql-parser", @@ -14098,7 +14041,6 @@ dependencies = [ "snafu 0.8.6", "sqlparser", "sqlparser_derive 0.1.1", - "store-api", "table", "uuid", ] @@ -14123,14 +14065,13 @@ dependencies = [ [[package]] name = "sqlness-runner" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "clap", "client", "common-error", "common-query", - "common-recordbatch", "common-time", "datatypes", "flate2", @@ -14403,7 +14344,7 @@ dependencies = [ [[package]] name = "standalone" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -14453,11 +14394,11 @@ checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" [[package]] name = "store-api" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "aquamarine", - "async-stream", + "arrow-schema 59.2.0", "async-trait", "bytes", "common-base", @@ -14484,7 +14425,6 @@ dependencies = [ "serde", "serde_json", "snafu 0.8.6", - "sqlparser", "strum 0.27.1", "tokio", "uuid", @@ -14646,23 +14586,20 @@ dependencies = [ [[package]] name = "substrait" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "bytes", "common-error", "common-function", "common-macro", - "common-telemetry", "datafusion", "datafusion-common", "datafusion-expr", "datafusion-substrait", - "datatypes", "promql", "prost 0.14.1", "snafu 0.8.6", - "tokio", ] [[package]] @@ -14785,7 +14722,7 @@ dependencies = [ [[package]] name = "table" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arc-swap", @@ -14811,7 +14748,6 @@ dependencies = [ "futures", "greptime-proto", "humantime", - "humantime-serde", "lazy_static", "once_cell", "parquet", @@ -14819,10 +14755,8 @@ dependencies = [ "serde", "serde_json", "snafu 0.8.6", - "sqlparser", "store-api", "tokio", - "tokio-util", ] [[package]] @@ -14936,9 +14870,9 @@ dependencies = [ [[package]] name = "tantivy-jieba" -version = "0.20.0" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3392170e86f1c387170aba7d171a466ffdc98a8b55b006e19ac64b123a7b690a" +checksum = "25cf6c5eb9419d7da9aacf5e184863b9ed476c478de0e196729f16124e716dd1" dependencies = [ "jieba-rs", "lazy_static", @@ -15068,7 +15002,7 @@ checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" [[package]] name = "tests-fuzz" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "arbitrary", "async-trait", @@ -15116,11 +15050,10 @@ dependencies = [ [[package]] name = "tests-integration" -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-flight", - "async-stream", "async-trait", "auth", "axum 0.8.4", @@ -15148,7 +15081,6 @@ dependencies = [ "common-query", "common-recordbatch", "common-runtime", - "common-stat", "common-telemetry", "common-test-util", "common-time", @@ -15179,7 +15111,6 @@ dependencies = [ "opentelemetry-proto", "operator", "otel-arrow-rust", - "partition", "paste", "pipeline", "plugins", @@ -15194,13 +15125,13 @@ dependencies = [ "servers", "session", "similar-asserts", - "snafu 0.8.6", "sql", "sqlx", "standalone", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.4.0-alpha.0", "table", + "temp-env", "tempfile", "time", "tokio", @@ -16226,16 +16157,6 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" -[[package]] -name = "usearch" -version = "2.21.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cc9fc5f872a3a4f9081d5f42624d788231b763e1846c829b9968a3755ac884d" -dependencies = [ - "cxx", - "cxx-build", -] - [[package]] name = "utf8-ranges" version = "1.0.5" @@ -16346,13 +16267,13 @@ version = "0.33.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "772206835b5bca8719825d2ab9e6fb160cbfce8ac1dc0e5dffe22cbe77254468" dependencies = [ - "aes", + "aes 0.8.4", "aes-siv", "base16", "base62", "base64-simd", "bytes", - "cbc", + "cbc 0.1.2", "cfb-mode", "cfg-if", "chacha20poly1305", @@ -16362,13 +16283,13 @@ dependencies = [ "ciborium", "cidr", "clap", - "codespan-reporting 0.12.0", + "codespan-reporting", "community-id", "convert_case 0.7.1", "crc", "crypto_secretbox", "csv", - "ctr", + "ctr 0.9.2", "digest 0.10.7", "dns-lookup", "domain", @@ -17429,9 +17350,9 @@ dependencies = [ [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "550a2b930b62486a393c52d5c3b84bff264b28aa437ed64694d31e93b1757af7" [[package]] name = "yaml-rust" diff --git a/Cargo.toml b/Cargo.toml index b4c4c9d7f54..cce9b502928 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,7 +76,7 @@ members = [ resolver = "2" [workspace.package] -version = "1.3.0-alpha.1" +version = "1.4.0-alpha.0" edition = "2024" license = "Apache-2.0" @@ -104,7 +104,6 @@ aquamarine = "0.6" arrow = { version = "=59.2.0", features = ["prettyprint"] } arrow-array = { version = "=59.2.0", default-features = false, features = ["chrono-tz"] } arrow-buffer = "=59.2.0" -arrow-cast = "=59.2.0" arrow-flight = "=59.2.0" arrow-ipc = { version = "=59.2.0", default-features = false, features = ["lz4", "zstd"] } arrow-schema = { version = "=59.2.0", features = ["serde"] } @@ -161,7 +160,8 @@ 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 = "0c4beffaa2b3968034915867593a717c6ab0c6b4" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" } +hdfs-native = "0.14.6" hex = "0.4" hostname = "0.4.0" http = "1" @@ -236,7 +236,6 @@ rskafka = { git = "https://github.com/influxdata/rskafka.git", rev = "80f42eb15e rstest = "0.25" rstest_reuse = "0.7" rust_decimal = "1.33" -rustc-hash = "2.0" rustls = { version = "0.23.25", default-features = false } ryu = "1.0" sea-query = "0.32" @@ -273,7 +272,6 @@ url = "2.3" # completion" panic (greptimedb#8265 / grpc-rust#2522). Remove once tonic ships a fix. tonic = { version = "0.14, <0.14.5", features = ["tls-aws-lc", "gzip", "zstd"] } tower = "0.5" -tower-http = "0.6" tracing = "0.1" tracing-appender = "0.2" tracing-opentelemetry = "0.33.0" @@ -365,21 +363,21 @@ rev = "f743c3a5fe8c1f57363eac5866c356b7d91f9b39" [patch.crates-io] substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" } -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } diff --git a/Makefile b/Makefile index 2526f3e31eb..7f92098262f 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ CARGO_BUILD_OPTS := --locked IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest -DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-04f07a91-20260921031255 +DEV_BUILDER_IMAGE_TAG ?= 1.98.1-7f0f592b-20261008091949 DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu @@ -212,12 +212,12 @@ sqlness-test: ## Run sqlness test. RUNS ?= 1 FUZZ_TARGET ?= fuzz_alter_table .PHONY: fuzz -fuzz: ## Run fuzz test ${FUZZ_TARGET}. - cargo fuzz run ${FUZZ_TARGET} --fuzz-dir tests-fuzz -D -s none -- -runs=${RUNS} +fuzz: ## Run fuzz test ${FUZZ_TARGET} (requires a nightly toolchain). + cargo +nightly fuzz run ${FUZZ_TARGET} --fuzz-dir tests-fuzz -D -s none -- -runs=${RUNS} .PHONY: fuzz-ls -fuzz-ls: ## List all fuzz targets. - cargo fuzz list --fuzz-dir tests-fuzz +fuzz-ls: ## List all fuzz targets (requires a nightly toolchain). + cargo +nightly fuzz list --fuzz-dir tests-fuzz .PHONY: check check: ## Cargo check all the targets. @@ -233,14 +233,11 @@ fix-clippy: ## Fix clippy violations. .PHONY: check-udeps check-udeps: ## Check unused dependencies. - cargo udeps --workspace --all-targets + cargo shear .PHONY: fix-udeps fix-udeps: ## Remove unused dependencies automatically. - @echo "Running cargo-udeps to find unused dependencies..." - @cargo udeps --workspace --all-targets --output json > udeps-report.json || true - @echo "Removing unused dependencies..." - @python3 scripts/fix-udeps.py udeps-report.json + cargo shear --fix .PHONY: fmt-check fmt-check: ## Check code format. @@ -264,6 +261,7 @@ stop-etcd: ## Stop single node etcd for testing purpose. run-it-in-container: start-etcd ## Run integration tests in dev-builder. docker run --network=host \ -v ${PWD}:/greptimedb -v ${CARGO_REGISTRY_CACHE}:/root/.cargo/registry -v ${CARGO_GIT_CACHE}:/root/.cargo/git -v /tmp:/tmp \ + -e GT_S3_BUCKET -e GT_S3_ACCESS_KEY_ID -e GT_S3_ACCESS_KEY -e GT_S3_REGION \ -w /greptimedb ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-${BASE_IMAGE}:${DEV_BUILDER_IMAGE_TAG} \ make test sqlness-test BUILD_JOBS=${BUILD_JOBS} diff --git a/README.md b/README.md index 1c8ed985e1d..96058cc79b7 100644 --- a/README.md +++ b/README.md @@ -168,7 +168,7 @@ Read more in the [full Install Guide](https://docs.greptime.com/getting-started/ ## Build From Source **Prerequisites:** -* [Rust toolchain](https://www.rust-lang.org/tools/install) — nightly, pinned by [`rust-toolchain.toml`](https://github.com/GreptimeTeam/greptimedb/blob/main/rust-toolchain.toml) +* [Rust toolchain](https://www.rust-lang.org/tools/install) — stable, pinned by [`rust-toolchain.toml`](https://github.com/GreptimeTeam/greptimedb/blob/main/rust-toolchain.toml) * [Protobuf compiler](https://grpc.io/docs/protoc-installation/) (>= 3.15) * C/C++ building essentials: `gcc` / `g++` / `autoconf` and the glibc dev package (`libc6-dev` on Ubuntu, `glibc-devel` on Fedora) * Python toolchain (optional, only for some test scripts) diff --git a/SECURITY.md b/SECURITY.md index ebb11e93c55..72029c7beb7 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -14,6 +14,9 @@ and cloud platform. If you come across a security vulnerability in GreptimeDB, we kindly request that you inform us immediately. We will thoroughly investigate all valid reports and make every effort to resolve the issue promptly. -To report any issues or vulnerabilities, please email us at info@greptime.com, rather than -posting publicly on GitHub. Be sure to provide us with the version identifier as well as details +Please do not report security vulnerabilities through public GitHub issues, discussions, or pull requests. + +Report them privately through GitHub's [private vulnerability reporting](https://github.com/GreptimeTeam/greptimedb/security/advisories/new) +(the "Report a vulnerability" button on the repository's Security tab). If you cannot use GitHub, +email info@greptime.com instead. Be sure to provide us with the version identifier as well as details on how the vulnerability can be exploited. diff --git a/config/config.md b/config/config.md index 16384b406f4..612a5cd43c0 100644 --- a/config/config.md +++ b/config/config.md @@ -16,17 +16,16 @@ | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | | `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | -| `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | -| `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL:
`static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | +| `max_in_flight_write_bytes` | String | 0 | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | +| `write_bytes_exhausted_policy` | String | `wait` | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `init_regions_in_background` | Bool | `false` | Initialize all regions in the background during the startup.
By default, it provides services after all regions have been initialized. | | `init_regions_parallelism` | Integer | `16` | Parallelism of initializing regions. | -| `max_concurrent_queries` | Integer | `0` | The maximum concurrent queries allowed to be executed. Zero means unlimited. | | `enable_telemetry` | Bool | `true` | Enable telemetry to collect anonymous usage data. Enabled by default. | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. | | `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency.
The weights set their relative runtime shares while both are backlogged. Disabled by default. | | `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. | @@ -34,7 +33,7 @@ | `runtime.experimental_workload_scheduler.sample_every_polls` | Integer | `16` | Number of polls between scheduler fairness samples. Must be greater than zero. | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | -| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | +| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | @@ -43,6 +42,8 @@ | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:4001` | The address to bind the gRPC server. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | +| `grpc.enable_cors` | Bool | `false` | Enable CORS for gRPC-Web clients in browsers. Disabled by default. | +| `grpc.cors_allowed_origins` | Array | Unset | Origins allowed by gRPC CORS. An empty list allows any origin. | | `grpc.max_connection_age` | String | Unset | The maximum connection age for gRPC connection.
The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
Refer to https://grpc.io/docs/guides/keepalive/ for more details. | | `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. | | `grpc.tls.mode` | String | `disable` | TLS mode. | @@ -75,33 +76,33 @@ | `influxdb` | -- | -- | InfluxDB protocol options. | | `influxdb.enable` | Bool | `true` | Whether to enable InfluxDB protocol in HTTP API. | | `influxdb.default_merge_mode` | String | `last_non_null` | Default merge mode for tables automatically created by InfluxDB protocol.
Available values: "last_non_null", "last_row". | -| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. | +| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in ingestion protocols.
PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge
queue admission without waiting for storage; later failures cannot be returned to the client.
Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.
OTLP logs, traces and ordinary metrics use this batcher.
MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged.
Single-connection writes and INSERT SELECT may incur additional flush waits. | | `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | -| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. | +| `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. | | `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | | `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | +| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | +| `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | +| `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | +| `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. | +| `pending_rows_batcher.logical_table.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per physical-table worker. | +| `pending_rows_batcher.logical_table.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | +| `pending_rows_batcher.logical_table.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued logical-table Flow notifications. | | `jaeger` | -- | -- | Jaeger protocol options. | | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | -| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.trace_aux_cache_size` | String | Auto | Estimated memory budget for cached trace service/operation keys per frontend,
shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching.
Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB.
Uses 32MiB if memory detection is unavailable.
Excludes cache and allocator overhead; shared table names are charged per entry. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | | `prom_store.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. | -| `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. | -| `prom_store.pending_rows_flush_interval` | String | `0s` | Interval to flush pending rows batcher.
Set to "0s" to disable batching mode in Prometheus Remote Write endpoint | -| `prom_store.max_batch_rows` | Integer | `100000` | Max rows per pending batch before triggering a flush. | -| `prom_store.max_concurrent_flushes` | Integer | `256` | Max number of concurrent batch flushes. | -| `prom_store.worker_channel_capacity` | Integer | `65526` | Capacity of the pending batch worker channel. | -| `prom_store.max_inflight_requests` | Integer | `3000` | Max inflight write requests before backpressure. | -| `prom_store.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of logical-table flow notifications waiting in the shared queue. | | `wal` | -- | -- | The WAL options. | -| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** | +| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** | | `wal.dir` | String | Unset | The directory to store the WAL files.
**It's only used when the provider is `raft_engine`**. | | `wal.file_size` | String | `128MB` | The size of the WAL segment file.
**It's only used when the provider is `raft_engine`**. | | `wal.purge_threshold` | String | `1GB` | The threshold of the WAL size to trigger a purge.
**It's only used when the provider is `raft_engine`**. | @@ -111,7 +112,7 @@ | `wal.enable_log_recycle` | Bool | `true` | Whether to reuse logically truncated log files.
**It's only used when the provider is `raft_engine`**. | | `wal.prefill_log_files` | Bool | `false` | Whether to pre-create log files on start up.
**It's only used when the provider is `raft_engine`**. | | `wal.sync_period` | String | `5s` | Duration for fsyncing log files.
**It's only used when the provider is `raft_engine`**. | -| `wal.recovery_parallelism` | Integer | `2` | Parallelism during WAL recovery. | +| `wal.recovery_parallelism` | Integer | Auto | Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. | | `wal.broker_endpoints` | Array | -- | The Kafka broker endpoints.
**It's only used when the provider is `kafka`**. | | `wal.connect_timeout` | String | `3s` | The connect timeout for kafka client.
**It's only used when the provider is `kafka`**. | | `wal.timeout` | String | `5s` | The total request timeout for kafka client.
**It's only used when the provider is `kafka`**. | @@ -126,7 +127,10 @@ | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `wal.storage_provider` | String | `""` | The name of the storage provider that holds the WAL objects, an empty name selects the default object store.
**It's only used when the provider is `experimental_object_store`**. | | `wal.prefix` | String | `wal` | The path prefix of the WAL objects inside the storage provider.
The objects are written under `/datanodes//epochs/`, which is derived from this prefix.
**It's only used when the provider is `experimental_object_store`**. | -| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.ack_mode` | String | `durable` | When an append to the object store WAL returns.
- `durable`: an append returns after the object holding its entries is durable (the default).
- `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | +| `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.
- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).
- `fail`: the read fails, so the region does not open.
**It's only used when the provider is `experimental_object_store`**. | | `metadata_store` | -- | -- | Metadata storage options. | | `metadata_store.file_size` | String | `64MB` | The size of the metadata store log file. | @@ -136,11 +140,9 @@ | `procedure.max_retry_times` | Integer | `3` | Procedure max retry time. | | `procedure.retry_delay` | String | `500ms` | Initial retry delay of procedures, increases exponentially | | `procedure.max_running_procedures` | Integer | `128` | Max running procedures.
The maximum number of procedures that can be running at the same time.
If the number of running procedures exceeds this limit, the procedure will be rejected. | -| `flow` | -- | -- | flow engine options. | -| `flow.num_workers` | Integer | `0` | The number of flow worker in flownode.
Not setting(or set to 0) this value will use the number of CPU cores divided by 2. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -149,9 +151,11 @@ | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | | `storage.copy_root` | String | `./greptimedb_data/copy` | Root directory for standalone SQL access to local files.
Relative SQL paths are resolved below this directory. Absolute paths are accepted only when
they are inside this directory. Defaults to `/copy`.
Distributed deployments always reject SQL access to local files.
Upgrade note: COPY commands and existing external tables that reference paths outside this
directory will fail. Move those files below the copy root, set this option to a dedicated
directory containing them, or migrate the files to object storage before upgrading. | -| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS. | +| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS.
- `Hdfs`: the data is stored in the Hadoop Distributed File System. | | `storage.bucket` | String | Unset | The S3 bucket name.
**It's only used when the storage type is `S3`, `Oss` and `Gcs`**. | -| `storage.root` | String | Unset | The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`.
**It's only used when the storage type is `S3`, `Oss` and `Azblob`**. | +| `storage.root` | String | Unset | The directory or object prefix under which data is stored.
**It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. | +| `storage.name_node` | String | Unset | The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`.
**It's only used when the storage type is `Hdfs`**. | +| `storage.options` | InlineTable | Unset | Additional options passed to the native HDFS client.
**It's only used when the storage type is `Hdfs`**. | | `storage.access_key_id` | String | Unset | The access key id of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3` and `Oss`**. | | `storage.secret_access_key` | String | Unset | The secret access key of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3`**. | | `storage.access_key_secret` | String | Unset | The secret access key of the aliyun account.
**It's only used when the storage type is `Oss`**. | @@ -172,12 +176,13 @@ | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | | `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | | `region_engine.mito` | -- | -- | The Mito engine options. | -| `region_engine.mito.num_workers` | Integer | `8` | Number of region workers. | +| `region_engine.mito.num_workers` | Integer | Auto | Number of region workers. Defaults to half of the available CPU cores, rounded up. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | | `region_engine.mito.worker_request_batch_size` | Integer | `64` | Max batch size for a worker to handle requests. | | `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_series_index_max_size` | String | `5GiB` | Approximate series and range index size limit in open regions, shared across workers.
Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation
can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots
retained by readers are not counted.
Minimum: 1KiB. Takes effect on restart. | | `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | | `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to
an exact multiple of each region's compaction time window. | @@ -212,7 +217,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

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

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | @@ -255,17 +260,18 @@ | `logging.tracing_sample_ratio` | -- | Unset | The percentage of tracing will be sampled and exported.
Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.
ratio > 1 are treated as 1. Fractions < 0 are treated as 0 | | `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- | | `slow_query` | -- | -- | The slow query log options. | -| `slow_query.enable` | Bool | `false` | Whether to enable slow query log. | -| `slow_query.record_type` | String | Unset | The record type of slow queries. It can be `system_table` or `log`. | -| `slow_query.threshold` | String | Unset | The threshold of slow query. | -| `slow_query.sample_ratio` | Float | Unset | The sampling ratio of slow query log. The value should be in the range of (0, 1]. | +| `slow_query.enable` | Bool | `true` | Whether to enable slow query log. | +| `slow_query.record_type` | String | `system_table` | The record type of slow queries. It can be `system_table` or `log`.
If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`.
If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. | +| `slow_query.threshold` | String | `30s` | The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. | +| `slow_query.sample_ratio` | Float | `1.0` | The sampling ratio of slow query log. The value should be in the range of (0, 1]. For example, `0.1` means 10% of the slow queries will be logged and `1.0` means all slow queries will be logged. | +| `slow_query.ttl` | String | `90d` | The TTL of the `slow_queries` system table. Default is `90d` when `record_type` is `system_table`. | | `tracing` | -- | -- | The tracing options. Only effect when compiled with `tokio-console` feature. | | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `pipeline` | -- | -- | The pipeline options. | | `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_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_recorder.event_types` | Array | All | 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. | @@ -280,16 +286,16 @@ | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | | `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | -| `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | -| `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL:
`static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | +| `max_in_flight_write_bytes` | String | 0 | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | +| `write_bytes_exhausted_policy` | String | `wait` | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | -| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | +| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | @@ -297,20 +303,22 @@ | `http.api_server_addr` | String | `127.0.0.1:4006` | The address to bind the dedicated HTTP API server, in the same form as `addr`.
Defaults to `127.0.0.1:4006`. | | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:4001` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:4001` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | -| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` | +| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `arrow_ipc`. | +| `grpc.enable_cors` | Bool | `false` | Enable CORS for gRPC-Web clients in browsers. Disabled by default. | +| `grpc.cors_allowed_origins` | Array | Unset | Origins allowed by gRPC CORS. An empty list allows any origin. | | `grpc.max_connection_age` | String | Unset | The maximum connection age for gRPC connection.
The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
Refer to https://grpc.io/docs/guides/keepalive/ for more details. | | `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. | | `grpc.tls.mode` | String | `disable` | TLS mode. | | `grpc.tls.cert_path` | String | Unset | Certificate file path. | | `grpc.tls.key_path` | String | Unset | Private key file path. | | `grpc.tls.watch` | Bool | `false` | Watch for Certificate and key file change and auto reload.
For now, gRPC tls config does not support auto reload. | -| `internal_grpc` | -- | -- | The internal gRPC server options. Internal gRPC port for nodes inside cluster to access frontend. | -| `internal_grpc.bind_addr` | String | `127.0.0.1:4010` | The address to bind the gRPC server. | -| `internal_grpc.server_addr` | String | `127.0.0.1:4010` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `internal_grpc.bind_addr`. | +| `internal_grpc` | -- | -- | The internal gRPC server options. Internal gRPC port for nodes inside cluster to access frontend.
The internal gRPC server starts only when this section is set. | +| `internal_grpc.bind_addr` | String | Unset | The address to bind the gRPC server.
If unset, it falls back to `127.0.0.1:4001`, the same address as `grpc.bind_addr`. | +| `internal_grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
Unlike `grpc.server_addr`, it is not detected automatically; set it explicitly. | | `internal_grpc.runtime_size` | Integer | `8` | The number of server worker threads. | -| `internal_grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` | +| `internal_grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `arrow_ipc`. | | `internal_grpc.tls` | -- | -- | internal gRPC server TLS options, see `mysql.tls` section. | | `internal_grpc.tls.mode` | String | `disable` | TLS mode. | | `internal_grpc.tls.cert_path` | String | Unset | Certificate file path. | @@ -342,31 +350,31 @@ | `influxdb` | -- | -- | InfluxDB protocol options. | | `influxdb.enable` | Bool | `true` | Whether to enable InfluxDB protocol in HTTP API. | | `influxdb.default_merge_mode` | String | `last_non_null` | Default merge mode for tables automatically created by InfluxDB protocol.
Available values: "last_non_null", "last_row". | -| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. | +| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in ingestion protocols.
PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge
queue admission without waiting for storage; later failures cannot be returned to the client.
Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.
OTLP logs, traces and ordinary metrics use this batcher.
MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged.
Single-connection writes and INSERT SELECT may incur additional flush waits. | | `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | -| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. | +| `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. | | `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | | `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | +| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | +| `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | +| `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | +| `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. | +| `pending_rows_batcher.logical_table.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per physical-table worker. | +| `pending_rows_batcher.logical_table.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | +| `pending_rows_batcher.logical_table.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued logical-table Flow notifications. | | `jaeger` | -- | -- | Jaeger protocol options. | | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | -| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.trace_aux_cache_size` | String | Auto | Estimated memory budget for cached trace service/operation keys per frontend,
shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching.
Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB.
Uses 32MiB if memory detection is unavailable.
Excludes cache and allocator overhead; shared table names are charged per entry. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | | `prom_store.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. | -| `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. | -| `prom_store.pending_rows_flush_interval` | String | `0s` | Interval to flush pending rows batcher.
Set to "0s" to disable batching mode in Prometheus Remote Write endpoint | -| `prom_store.max_batch_rows` | Integer | `100000` | Max rows per pending batch before triggering a flush. | -| `prom_store.max_concurrent_flushes` | Integer | `256` | Max number of concurrent batch flushes. | -| `prom_store.worker_channel_capacity` | Integer | `65526` | Capacity of the pending batch worker channel. | -| `prom_store.max_inflight_requests` | Integer | `3000` | Max inflight write requests before backpressure. | -| `prom_store.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of logical-table flow notifications waiting in the shared queue. | | `meta_client` | -- | -- | The metasrv client options. | | `meta_client.metasrv_addrs` | Array | -- | The addresses of the metasrv. | | `meta_client.timeout` | String | `3s` | Operation timeout. | @@ -379,7 +387,7 @@ | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | | `query.allow_query_fallback` | Bool | `false` | Whether to allow query fallback when push down optimize fails.
Default to false, meaning when push down optimize failed, return error msg | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans (only applies to datanodes). | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans (only applies to datanodes). | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -387,7 +395,7 @@ | `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `datanode` | -- | -- | Datanode options. | | `datanode.client` | -- | -- | Datanode client options. | -| `datanode.client.connect_timeout` | String | `10s` | -- | +| `datanode.client.connect_timeout` | String | 1s | Connect server timeout. Defaults to `1s`; this example sets it to `10s`. | | `datanode.client.tcp_nodelay` | Bool | `true` | -- | | `datanode.client.max_recv_message_size` | String | `512MB` | The maximum receive message size for the gRPC client. | | `datanode.client.max_send_message_size` | String | `512MB` | The maximum send message size for the gRPC client. | @@ -420,7 +428,7 @@ | `pipeline.cache_ttl` | String | `10s` | Time to live of the frontend-local pipeline cache. A pipeline created or deleted on
another frontend takes effect on this one after at most this duration. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | All | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | ### Metasrv @@ -432,7 +440,7 @@ | `store_key_prefix` | String | `""` | If it's not empty, the metasrv will store all data with this key prefix. | | `backend` | String | `etcd_store` | The datastore for meta server.
Available values:
- `etcd_store` (default value)
- `memory_store`
- `postgres_store`
- `mysql_store` | | `meta_table_name` | String | `greptime_metakv` | Table name in RDS to store metadata. Effect when using a RDS kvbackend.
**Only used when backend is `postgres_store`.** | -| `meta_schema_name` | String | `greptime_schema` | Optional PostgreSQL schema for metadata table and election table name qualification.
When PostgreSQL public schema is not writable (e.g., PostgreSQL 15+ with restricted public),
set this to a writable schema. GreptimeDB will use `meta_schema_name`.`meta_table_name`.
**Only used when backend is `postgres_store`.** | +| `meta_schema_name` | String | Unset | Optional PostgreSQL schema for metadata table and election table name qualification.
When PostgreSQL public schema is not writable (e.g., PostgreSQL 15+ with restricted public),
set this to a writable schema. GreptimeDB will use `meta_schema_name`.`meta_table_name`.
**Only used when backend is `postgres_store`.**
Unset by default; this example sets it to `greptime_schema`. | | `auto_create_schema` | Bool | `true` | Automatically create PostgreSQL schema if it doesn't exist.
When enabled, the system will execute `CREATE SCHEMA IF NOT EXISTS `
before creating metadata tables. This is useful in production environments where
manual schema creation may be restricted.
Default is true.
Note: The PostgreSQL user must have CREATE SCHEMA permission for this to work.
**Only used when backend is `postgres_store`.** | | `meta_election_lock_id` | Integer | `1` | Advisory lock id in PostgreSQL for election. Effect when using PostgreSQL as kvbackend
Only used when backend is `postgres_store`. | | `selector` | String | `round_robin` | Datanode selector type.
- `round_robin` (default value)
- `lease_based`
- `load_based`
For details, please see "https://docs.greptime.com/developer-guide/metasrv/selector". | @@ -443,9 +451,9 @@ | `heartbeat_interval` | String | `3s` | Base heartbeat interval for calculating distributed time constants.
The frontend heartbeat interval is 6 times of the base heartbeat interval.
The flownode/datanode heartbeat interval is 1 times of the base heartbeat interval.
e.g., If the base heartbeat interval is 3s, the frontend heartbeat interval is 18s, the flownode/datanode heartbeat interval is 3s.
Heartbeat intervals are negotiated from metasrv during handshake; local node configs do not override this. | | `enable_telemetry` | Bool | `true` | Whether to enable greptimedb telemetry. Enabled by default. | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `backend_tls` | -- | -- | TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends)
When using etcd, PostgreSQL, or MySQL as metadata store, you can configure TLS here

Note: if TLS is configured in both this section and the `store_addrs` connection string, the
settings here will override the TLS settings in `store_addrs`. | | `backend_tls.mode` | String | `prefer` | TLS mode, refer to https://www.postgresql.org/docs/current/libpq-ssl.html
- "disable" - No TLS
- "prefer" (default) - Try TLS, fallback to plain
- "require" - Require TLS
- "verify_ca" - Require TLS and verify CA
- "verify_full" - Require TLS and verify hostname | | `backend_tls.cert_path` | String | `""` | Path to client certificate file (for client authentication)
Like "/path/to/client.crt" | @@ -457,7 +465,7 @@ | `backend_client.connect_timeout` | String | `3s` | The connect timeout for backend client. | | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:3002` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:3002` | The communication server address for the frontend and datanode to connect to metasrv.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.server_addr` | String | Unset | The communication server address for the frontend and datanode to connect to metasrv.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | | `grpc.max_recv_message_size` | String | `512MB` | The maximum receive message size for gRPC server. | | `grpc.max_send_message_size` | String | `512MB` | The maximum send message size for gRPC server. | @@ -479,7 +487,7 @@ | `datanode` | -- | -- | Datanode options. | | `datanode.client` | -- | -- | Datanode client options. | | `datanode.client.timeout` | String | `10s` | Operation timeout. | -| `datanode.client.connect_timeout` | String | `10s` | Connect server timeout. | +| `datanode.client.connect_timeout` | String | 1s | Connect server timeout. Defaults to `1s`; this example sets it to `10s`. | | `datanode.client.tcp_nodelay` | Bool | `true` | `TCP_NODELAY` option for accepted connections. | | `datanode.client.max_recv_message_size` | String | `512MB` | The maximum receive message size for the gRPC client. | | `datanode.client.max_send_message_size` | String | `512MB` | The maximum send message size for the gRPC client. | @@ -501,7 +509,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_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_recorder.event_types` | Array | All | 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`. | @@ -546,22 +554,22 @@ | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:3001` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:3001` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | | `grpc.max_recv_message_size` | String | `512MB` | The maximum receive message size for gRPC server. | | `grpc.max_send_message_size` | String | `512MB` | The maximum send message size for gRPC server. | -| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for datanode side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` | +| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for datanode side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `arrow_ipc`. | | `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. | | `grpc.tls.mode` | String | `disable` | TLS mode. | | `grpc.tls.cert_path` | String | Unset | Certificate file path. | | `grpc.tls.key_path` | String | Unset | Private key file path. | | `grpc.tls.watch` | Bool | `false` | Watch for Certificate and key file change and auto reload.
For now, gRPC tls config does not support auto reload. | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | -| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 2). | -| `runtime.ingest_rt_size` | Integer | `8` | The number of threads to execute datanode ingestion operations. | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.query_rt_size` | Integer | Auto | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 2). | +| `runtime.ingest_rt_size` | Integer | Auto | The number of threads to execute datanode ingestion operations.
Defaults to max(num_cpus, 2). | | `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. | | `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency.
The weights set their relative runtime shares while both are backlogged. Disabled by default. | | `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. | @@ -577,7 +585,7 @@ | `meta_client.metadata_cache_ttl` | String | `10m` | TTL of the metadata cache. | | `meta_client.metadata_cache_tti` | String | `5m` | -- | | `wal` | -- | -- | The WAL options. | -| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `noop`: it's a no-op WAL provider that does not store any WAL data.
**Notes: any unflushed data will be lost when the datanode is shutdown.**
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** | +| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `noop`: it's a no-op WAL provider that does not store any WAL data.
**Notes: any unflushed data will be lost when the datanode is shutdown.**
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** | | `wal.dir` | String | Unset | The directory to store the WAL files.
**It's only used when the provider is `raft_engine`**. | | `wal.file_size` | String | `128MB` | The size of the WAL segment file.
**It's only used when the provider is `raft_engine`**. | | `wal.purge_threshold` | String | `1GB` | The threshold of the WAL size to trigger a purge.
**It's only used when the provider is `raft_engine`**. | @@ -587,7 +595,7 @@ | `wal.enable_log_recycle` | Bool | `true` | Whether to reuse logically truncated log files.
**It's only used when the provider is `raft_engine`**. | | `wal.prefill_log_files` | Bool | `false` | Whether to pre-create log files on start up.
**It's only used when the provider is `raft_engine`**. | | `wal.sync_period` | String | `5s` | Duration for fsyncing log files.
**It's only used when the provider is `raft_engine`**. | -| `wal.recovery_parallelism` | Integer | `2` | Parallelism during WAL recovery. | +| `wal.recovery_parallelism` | Integer | Auto | Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. | | `wal.broker_endpoints` | Array | -- | The Kafka broker endpoints.
**It's only used when the provider is `kafka`**. | | `wal.connect_timeout` | String | `3s` | The connect timeout for kafka client.
**It's only used when the provider is `kafka`**. | | `wal.timeout` | String | `5s` | The total request timeout for kafka client.
**It's only used when the provider is `kafka`**. | @@ -598,11 +606,14 @@ | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `wal.storage_provider` | String | `""` | The name of the storage provider that holds the WAL objects, an empty name selects the default object store.
**It's only used when the provider is `experimental_object_store`**. | | `wal.prefix` | String | `wal` | The path prefix of the WAL objects inside the storage provider.
The objects are written under `/datanodes//epochs/`, which is derived from this prefix.
**It's only used when the provider is `experimental_object_store`**. | -| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.ack_mode` | String | `durable` | When an append to the object store WAL returns.
- `durable`: an append returns after the object holding its entries is durable (the default).
- `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | +| `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.
- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).
- `fail`: the read fails, so the region does not open.
**It's only used when the provider is `experimental_object_store`**. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -610,9 +621,11 @@ | `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | -| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS. | +| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS.
- `Hdfs`: the data is stored in the Hadoop Distributed File System. | | `storage.bucket` | String | Unset | The S3 bucket name.
**It's only used when the storage type is `S3`, `Oss` and `Gcs`**. | -| `storage.root` | String | Unset | The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`.
**It's only used when the storage type is `S3`, `Oss` and `Azblob`**. | +| `storage.root` | String | Unset | The directory or object prefix under which data is stored.
**It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. | +| `storage.name_node` | String | Unset | The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`.
**It's only used when the storage type is `Hdfs`**. | +| `storage.options` | InlineTable | Unset | Additional options passed to the native HDFS client.
**It's only used when the storage type is `Hdfs`**. | | `storage.access_key_id` | String | Unset | The access key id of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3` and `Oss`**. | | `storage.secret_access_key` | String | Unset | The secret access key of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3`**. | | `storage.access_key_secret` | String | Unset | The secret access key of the aliyun account.
**It's only used when the storage type is `Oss`**. | @@ -633,7 +646,7 @@ | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | | `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | | `region_engine.mito` | -- | -- | The Mito engine options. | -| `region_engine.mito.num_workers` | Integer | `8` | Number of region workers. | +| `region_engine.mito.num_workers` | Integer | Auto | Number of region workers. Defaults to half of the available CPU cores, rounded up. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | | `region_engine.mito.worker_request_batch_size` | Integer | `64` | Max batch size for a worker to handle requests. | | `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. | @@ -641,6 +654,7 @@ | `region_engine.mito.experimental_manifest_keep_removed_file_ttl` | String | `1h` | How long to keep removed files in the `removed_files` field of manifest
after they are removed from manifest.
files will only be removed from `removed_files` field
if both `keep_removed_file_count` and `keep_removed_file_ttl` is reached. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_series_index_max_size` | String | `5GiB` | Approximate series and range index size limit in open regions, shared across workers.
Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation
can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots
retained by readers are not counted.
Minimum: 1KiB. Takes effect on restart. | | `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | | `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to
an exact multiple of each region's compaction time window. | @@ -675,7 +689,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

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

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | @@ -732,7 +746,6 @@ | --- | -----| ------- | ----------- | | `node_id` | Integer | Unset | The flownode identifier and should be unique in the cluster. | | `flow` | -- | -- | flow engine options. | -| `flow.num_workers` | Integer | `0` | The number of flow worker in flownode.
Not setting(or set to 0) this value will use the number of CPU cores divided by 2. | | `flow.batching_mode` | -- | -- | -- | | `flow.batching_mode.query_timeout` | String | `600s` | The default batching engine query timeout is 10 minutes. | | `flow.batching_mode.slow_query_threshold` | String | `60s` | will output a warn log for any query that runs for more that this threshold | @@ -750,9 +763,9 @@ | `flow.batching_mode.frontend_tls.client_cert_path` | String | Unset | Client Certificate file path. | | `flow.batching_mode.frontend_tls.client_key_path` | String | Unset | Client Private key file path. | | `grpc` | -- | -- | The gRPC server options. | -| `grpc.bind_addr` | String | `127.0.0.1:6800` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:6800` | The address advertised to the metasrv,
and used for connections from outside the host | -| `grpc.runtime_size` | Integer | `2` | The number of server worker threads. | +| `grpc.bind_addr` | String | 127.0.0.1:3004 | The address to bind the gRPC server. Defaults to `127.0.0.1:3004`; this example sets it to `127.0.0.1:6800`. | +| `grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.runtime_size` | Integer | 8 | The number of server worker threads. Defaults to 8; this example sets it to 2. | | `grpc.max_recv_message_size` | String | `512MB` | The maximum receive message size for gRPC server. | | `grpc.max_send_message_size` | String | `512MB` | The maximum send message size for gRPC server. | | `http` | -- | -- | The HTTP server options. | @@ -786,7 +799,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `query` | -- | -- | -- | | `query.parallelism` | Integer | `1` | Parallelism of the query engine for query sent by flownode.
Default to 1, so it won't use too much cpu or memory | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 45baaefbdfb..413afb609c6 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -45,6 +45,7 @@ bind_addr = "127.0.0.1:3001" ## The address advertised to the metasrv, and used for connections from outside the host. ## If left empty or unset, the server will automatically use the IP address of the first network interface ## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:3001" ## The number of server worker threads. runtime_size = 8 @@ -57,7 +58,7 @@ max_send_message_size = "512MB" ## - `transport`: only enable gRPC transport compression (zstd) ## - `arrow_ipc`: only enable Arrow IPC compression (lz4) ## - `all`: enable all compression. -## Default to `none` +## Default to `arrow_ipc`. flight_compression = "arrow_ipc" ## gRPC server TLS options, see `mysql.tls` section. @@ -80,16 +81,24 @@ watch = false ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## The number of threads to execute datanode query operations. ## Defaults to max(num_cpus - 1, 2). +## @toml2docs:none-default="Auto" #+ query_rt_size = 7 ## The number of threads to execute datanode ingestion operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ ingest_rt_size = 8 ## Experimental weighted, work-conserving query/write task scheduler. @@ -136,7 +145,7 @@ metadata_cache_tti = "5m" ## - `raft_engine`: the wal is stored in the local file system by raft-engine. ## - `kafka`: it's remote wal that data is stored in Kafka. ## - `noop`: it's a no-op WAL provider that does not store any WAL data.
**Notes: any unflushed data will be lost when the datanode is shutdown.** -## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** +## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** provider = "raft_engine" ## The directory to store the WAL files. @@ -176,7 +185,8 @@ prefill_log_files = false ## **It's only used when the provider is `raft_engine`**. sync_period = "5s" -## Parallelism during WAL recovery. +## Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. +## @toml2docs:none-default="Auto" recovery_parallelism = 2 ## The Kafka broker endpoints. @@ -231,10 +241,25 @@ overwrite_entry_start_id = false ## The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`. ## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters. -## When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node. +## When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission. ## **It's only used when the provider is `experimental_object_store`**. #+ flush_interval = "100ms" +## When an append to the object store WAL returns. +## - `durable`: an append returns after the object holding its entries is durable (the default). +## - `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog. +## **It's only used when the provider is `experimental_object_store`**. +#+ ack_mode = "durable" + +## The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode. +## Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_bytes = "64MB" + +## The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_age = "8s" + ## What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry. ## - `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default). ## - `fail`: the read fails, so the region does not open. @@ -300,6 +325,13 @@ overwrite_entry_start_id = false # credential = "base64-credential" # endpoint = "https://storage.googleapis.com" +# Example of using HDFS as the storage with the native Rust client. +# [storage] +# type = "Hdfs" +# root = "/greptimedb" +# name_node = "hdfs://127.0.0.1:9000" +# options = { "dfs.client.block.write.replace-datanode-on-failure.enable" = "true" } + ## The query engine options. [query] ## Parallelism of the query engine. @@ -308,9 +340,10 @@ parallelism = 0 ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans. +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: @@ -348,6 +381,7 @@ data_home = "./greptimedb_data" ## - `Gcs`: the data is stored in the Google Cloud Storage. ## - `Azblob`: the data is stored in the Azure Blob Storage. ## - `Oss`: the data is stored in the Aliyun OSS. +## - `Hdfs`: the data is stored in the Hadoop Distributed File System. type = "File" ## The S3 bucket name. @@ -355,11 +389,21 @@ type = "File" ## @toml2docs:none-default bucket = "greptimedb" -## The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`. -## **It's only used when the storage type is `S3`, `Oss` and `Azblob`**. +## The directory or object prefix under which data is stored. +## **It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. ## @toml2docs:none-default root = "greptimedb" +## The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +name_node = "hdfs://127.0.0.1:9000" + +## Additional options passed to the native HDFS client. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +options = {} + ## The access key id of the aws account. ## It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key. ## **It's only used when the storage type is `S3` and `Oss`**. @@ -470,7 +514,8 @@ skip_ssl_validation = false ## The Mito engine options. [region_engine.mito] -## Number of region workers. +## Number of region workers. Defaults to half of the available CPU cores, rounded up. +## @toml2docs:none-default="Auto" #+ num_workers = 8 ## Request channel size of each worker. @@ -502,6 +547,13 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Approximate series and range index size limit in open regions, shared across workers. +## Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation +## can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots +## retained by readers are not counted. +## Minimum: 1KiB. Takes effect on restart. +#+ experimental_series_index_max_size = "5GiB" + ## Whether to build and query range indexes when series indexes are enabled. ## Obsolete range-index metadata and files are still cleaned up when disabled. #+ experimental_enable_range_index = false @@ -647,7 +699,7 @@ schedule_compaction_after_edit = true default_flat_format = true ## Whether to enable the experimental two-phase mode for eligible metric series scans. -experimental_series_scan_v2 = true +experimental_series_scan_v2 = false ## The options for index in Mito engine. [region_engine.mito.index] diff --git a/config/flownode.example.toml b/config/flownode.example.toml index 192c6af6977..fab31fed510 100644 --- a/config/flownode.example.toml +++ b/config/flownode.example.toml @@ -4,9 +4,6 @@ node_id = 14 ## flow engine options. [flow] -## The number of flow worker in flownode. -## Not setting(or set to 0) this value will use the number of CPU cores divided by 2. -#+num_workers=0 [flow.batching_mode] ## The default batching engine query timeout is 10 minutes. #+query_timeout="600s" @@ -47,12 +44,16 @@ node_id = 14 ## The gRPC server options. [grpc] -## The address to bind the gRPC server. +## The address to bind the gRPC server. Defaults to `127.0.0.1:3004`; this example sets it to `127.0.0.1:6800`. +## @toml2docs:none-default="127.0.0.1:3004" bind_addr = "127.0.0.1:6800" -## The address advertised to the metasrv, -## and used for connections from outside the host +## The address advertised to the metasrv, and used for connections from outside the host. +## If left empty or unset, the server will automatically use the IP address of the first network interface +## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:6800" -## The number of server worker threads. +## The number of server worker threads. Defaults to 8; this example sets it to 2. +## @toml2docs:none-default="8" runtime_size = 2 ## The maximum receive message size for gRPC server. max_recv_message_size = "512MB" @@ -156,9 +157,10 @@ parallelism = 1 ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans. +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 043da57db74..e4ef6e36935 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -20,8 +20,8 @@ default_column_prefix = "greptime" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", ## "mysql_native_password:", ## "pg_scram_sha256::::" -## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's -## native password handshake; a MySQL client must send the password in cleartext for such users. +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL: +## `static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats. ## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. ## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or ## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. @@ -33,22 +33,26 @@ default_column_prefix = "greptime" ## Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). ## Set to 0 to disable the limit. Default: "0" (unlimited) -## @toml2docs:none-default +## @toml2docs:none-default="0" #+ max_in_flight_write_bytes = "1GB" ## Policy when write bytes quota is exhausted. ## Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" -## @toml2docs:none-default #+ write_bytes_exhausted_policy = "wait" ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## The HTTP server options. @@ -56,9 +60,9 @@ default_column_prefix = "greptime" ## The address to bind the HTTP server. addr = "127.0.0.1:4000" ## HTTP request timeout. Set to 0 to disable timeout. -## When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is +## When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is ## raised to at least the largest active flush interval plus 1 second. The intervals come from -## `prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. +## `pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. timeout = "0s" ## HTTP request body limit. ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`. @@ -86,6 +90,7 @@ bind_addr = "127.0.0.1:4001" ## The address advertised to the metasrv, and used for connections from outside the host. ## If left empty or unset, the server will automatically use the IP address of the first network interface ## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:4001" ## The number of server worker threads. runtime_size = 8 @@ -94,8 +99,13 @@ runtime_size = 8 ## - `transport`: only enable gRPC transport compression (zstd) ## - `arrow_ipc`: only enable Arrow IPC compression (lz4) ## - `all`: enable all compression. -## Default to `none` +## Default to `arrow_ipc`. flight_compression = "arrow_ipc" +## Enable CORS for gRPC-Web clients in browsers. Disabled by default. +#+ enable_cors = false +## Origins allowed by gRPC CORS. An empty list allows any origin. +## @toml2docs:none-default +#+ cors_allowed_origins = ["https://example.com"] ## The maximum connection age for gRPC connection. ## The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour. ## Refer to https://grpc.io/docs/guides/keepalive/ for more details. @@ -120,12 +130,15 @@ key_path = "" watch = false ## The internal gRPC server options. Internal gRPC port for nodes inside cluster to access frontend. +## The internal gRPC server starts only when this section is set. [internal_grpc] ## The address to bind the gRPC server. +## If unset, it falls back to `127.0.0.1:4001`, the same address as `grpc.bind_addr`. +## @toml2docs:none-default bind_addr = "127.0.0.1:4010" ## The address advertised to the metasrv, and used for connections from outside the host. -## If left empty or unset, the server will automatically use the IP address of the first network interface -## on the host, with the same port number as the one specified in `internal_grpc.bind_addr`. +## Unlike `grpc.server_addr`, it is not detected automatically; set it explicitly. +## @toml2docs:none-default server_addr = "127.0.0.1:4010" ## The number of server worker threads. runtime_size = 8 @@ -134,7 +147,7 @@ runtime_size = 8 ## - `transport`: only enable gRPC transport compression (zstd) ## - `arrow_ipc`: only enable Arrow IPC compression (lz4) ## - `all`: enable all compression. -## Default to `none` +## Default to `arrow_ipc`. flight_compression = "arrow_ipc" ## internal gRPC server TLS options, see `mysql.tls` section. @@ -231,12 +244,13 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" -## Shared experimental ordinary-table batching for opted-in ingestion protocols. -## Legacy Prometheus batching settings under prom_store remain supported. -## HTTP write protocols sharing this batcher. Omitted or empty disables all entrances. -## Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom. -## Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher. -## Effective shared Prom settings take precedence; existing prom_store settings remain compatible. +## Ordinary-table batching for opted-in ingestion protocols. +## PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge +## queue admission without waiting for storage; later failures cannot be returned to the client. +## Omitted or empty protocols disables batching. Prom without metric engine uses this batcher. +## OTLP logs, traces and ordinary metrics use this batcher. +## MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged. +## Single-connection writes and INSERT SELECT may incur additional flush waits. [pending_rows_batcher] # protocols = [ # "influxdb", @@ -247,6 +261,8 @@ default_merge_mode = "last_non_null" # "splunk", # "elasticsearch", # "http_sql", +# "mysql", +# "postgres", # "prom", # ] ## Flush interval measured from the first pending submission. Zero disables batching. @@ -256,12 +272,33 @@ max_batch_rows = 100000 ## Maximum concurrent flushes shared by the frontend batcher. max_concurrent_flushes = 256 ## Maximum queued submissions per table worker. -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 ## Maximum admitted original requests awaiting completion. max_inflight_requests = 3000 ## Maximum number of queued table Flow notifications. flow_notification_queue_capacity = 1024 +## Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics. +## Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible. +## Enable independently with protocols and a nonzero flush interval. +## Omitted fields use independent defaults, not parent settings. +## Empty protocols or a zero interval disables logical batching without fallback. +## Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. +[pending_rows_batcher.logical_table] +# protocols = ["prom", "otlp"] +## Flush interval measured from the first pending submission. Zero disables batching. +pending_rows_flush_interval = "0s" +## Flush after a complete submission reaches this row threshold. +max_batch_rows = 100000 +## Maximum concurrent flushes shared by Prom and OTLP metrics. +max_concurrent_flushes = 256 +## Maximum queued submissions per physical-table worker. +worker_channel_capacity = 65536 +## Maximum admitted original requests awaiting completion. +max_inflight_requests = 3000 +## Maximum number of queued logical-table Flow notifications. +flow_notification_queue_capacity = 1024 + ## Jaeger protocol options. [jaeger] ## Whether to enable Jaeger protocol in HTTP API. @@ -271,10 +308,15 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true -## Experimental: enable cumulative OTLP exponential histogram ingestion. -experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Estimated memory budget for cached trace service/operation keys per frontend, +## shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. +## Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. +## Uses 32MiB if memory detection is unavailable. +## Excludes cache and allocator overhead; shared table names are charged per entry. +## @toml2docs:none-default="Auto" +#+ trace_aux_cache_size = "32MiB" ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric ## resource attributes, so metrics-only services reach the semantic graph. experimental_enable_resource_info = false @@ -291,21 +333,7 @@ with_metric_engine = true ## - lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD). ## - unchecked: do not valid strings. prom_validation_mode = "strict" -## Experimental: enable Prometheus remote write v2 native histogram ingestion. -experimental_enable_prometheus_native_histogram = false -## Interval to flush pending rows batcher. -## Set to "0s" to disable batching mode in Prometheus Remote Write endpoint -#+pending_rows_flush_interval = "0s" -## Max rows per pending batch before triggering a flush. -#+max_batch_rows = 100000 -## Max number of concurrent batch flushes. -#+max_concurrent_flushes = 256 -## Capacity of the pending batch worker channel. -#+worker_channel_capacity = 65526 -## Max inflight write requests before backpressure. -#+max_inflight_requests = 3000 -## Maximum number of logical-table flow notifications waiting in the shared queue. -#+flow_notification_queue_capacity = 1024 + ## The metasrv client options. [meta_client] @@ -344,9 +372,10 @@ allow_query_fallback = false ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans (only applies to datanodes). +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: @@ -377,6 +406,8 @@ memory_pool_size = "50%" [datanode] ## Datanode client options. [datanode.client] +## Connect server timeout. Defaults to `1s`; this example sets it to `10s`. +## @toml2docs:none-default="1s" connect_timeout = "10s" tcp_nodelay = true ## The maximum receive message size for the gRPC client. @@ -481,4 +512,5 @@ ttl = "90d" ## Event types to record. Current available event type: `admin_function`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. +## @toml2docs:none-default="All" #+ event_types = ["admin_function"] diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index b923c5a549c..fe4c38d02f9 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -35,6 +35,8 @@ meta_table_name = "greptime_metakv" ## When PostgreSQL public schema is not writable (e.g., PostgreSQL 15+ with restricted public), ## set this to a writable schema. GreptimeDB will use `meta_schema_name`.`meta_table_name`. ## **Only used when backend is `postgres_store`.** +## Unset by default; this example sets it to `greptime_schema`. +## @toml2docs:none-default meta_schema_name = "greptime_schema" ## Automatically create PostgreSQL schema if it doesn't exist. @@ -88,11 +90,16 @@ node_max_idle_time = "24hours" ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends) @@ -140,6 +147,7 @@ bind_addr = "127.0.0.1:3002" ## The communication server address for the frontend and datanode to connect to metasrv. ## If left empty or unset, the server will automatically use the IP address of the first network interface ## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:3002" ## The number of server worker threads. runtime_size = 8 @@ -206,7 +214,8 @@ acceptable_heartbeat_pause = "10000ms" ## Operation timeout. timeout = "10s" -## Connect server timeout. +## Connect server timeout. Defaults to `1s`; this example sets it to `10s`. +## @toml2docs:none-default="1s" connect_timeout = "10s" ## `TCP_NODELAY` option for accepted connections. @@ -332,6 +341,7 @@ ttl = "90d" ## `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. +## @toml2docs:none-default="All" #+ event_types = ["region_migration"] ## Configuration options for the stats persistence. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index e59e7ba1b4c..c656049c5d2 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -20,8 +20,8 @@ default_column_prefix = "greptime" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", ## "mysql_native_password:", ## "pg_scram_sha256::::" -## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's -## native password handshake; a MySQL client must send the password in cleartext for such users. +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL: +## `static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats. ## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. ## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or ## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. @@ -33,12 +33,11 @@ default_column_prefix = "greptime" ## Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). ## Set to 0 to disable the limit. Default: "0" (unlimited) -## @toml2docs:none-default +## @toml2docs:none-default="0" #+ max_in_flight_write_bytes = "1GB" ## Policy when write bytes quota is exhausted. ## Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" -## @toml2docs:none-default #+ write_bytes_exhausted_policy = "wait" ## Initialize all regions in the background during the startup. @@ -48,20 +47,22 @@ init_regions_in_background = false ## Parallelism of initializing regions. init_regions_parallelism = 16 -## The maximum concurrent queries allowed to be executed. Zero means unlimited. -max_concurrent_queries = 0 - ## Enable telemetry to collect anonymous usage data. Enabled by default. #+ enable_telemetry = true ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## Experimental weighted, work-conserving query/write task scheduler. @@ -81,9 +82,9 @@ max_concurrent_queries = 0 ## The address to bind the HTTP server. addr = "127.0.0.1:4000" ## HTTP request timeout. Set to 0 to disable timeout. -## When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is +## When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is ## raised to at least the largest active flush interval plus 1 second. The intervals come from -## `prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. +## `pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. timeout = "0s" ## HTTP request body limit. ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`. @@ -111,6 +112,11 @@ api_server_addr = "127.0.0.1:4006" bind_addr = "127.0.0.1:4001" ## The number of server worker threads. runtime_size = 8 +## Enable CORS for gRPC-Web clients in browsers. Disabled by default. +#+ enable_cors = false +## Origins allowed by gRPC CORS. An empty list allows any origin. +## @toml2docs:none-default +#+ cors_allowed_origins = ["https://example.com"] ## The maximum connection age for gRPC connection. ## The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour. ## Refer to https://grpc.io/docs/guides/keepalive/ for more details. @@ -210,12 +216,13 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" -## Shared experimental ordinary-table batching for opted-in ingestion protocols. -## Legacy Prometheus batching settings under prom_store remain supported. -## HTTP write protocols sharing this batcher. Omitted or empty disables all entrances. -## Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom. -## Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher. -## Effective shared Prom settings take precedence; existing prom_store settings remain compatible. +## Ordinary-table batching for opted-in ingestion protocols. +## PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge +## queue admission without waiting for storage; later failures cannot be returned to the client. +## Omitted or empty protocols disables batching. Prom without metric engine uses this batcher. +## OTLP logs, traces and ordinary metrics use this batcher. +## MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged. +## Single-connection writes and INSERT SELECT may incur additional flush waits. [pending_rows_batcher] # protocols = [ # "influxdb", @@ -226,6 +233,8 @@ default_merge_mode = "last_non_null" # "splunk", # "elasticsearch", # "http_sql", +# "mysql", +# "postgres", # "prom", # ] ## Flush interval measured from the first pending submission. Zero disables batching. @@ -235,12 +244,33 @@ max_batch_rows = 100000 ## Maximum concurrent flushes shared by the frontend batcher. max_concurrent_flushes = 256 ## Maximum queued submissions per table worker. -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 ## Maximum admitted original requests awaiting completion. max_inflight_requests = 3000 ## Maximum number of queued table Flow notifications. flow_notification_queue_capacity = 1024 +## Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics. +## Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible. +## Enable independently with protocols and a nonzero flush interval. +## Omitted fields use independent defaults, not parent settings. +## Empty protocols or a zero interval disables logical batching without fallback. +## Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. +[pending_rows_batcher.logical_table] +# protocols = ["prom", "otlp"] +## Flush interval measured from the first pending submission. Zero disables batching. +pending_rows_flush_interval = "0s" +## Flush after a complete submission reaches this row threshold. +max_batch_rows = 100000 +## Maximum concurrent flushes shared by Prom and OTLP metrics. +max_concurrent_flushes = 256 +## Maximum queued submissions per physical-table worker. +worker_channel_capacity = 65536 +## Maximum admitted original requests awaiting completion. +max_inflight_requests = 3000 +## Maximum number of queued logical-table Flow notifications. +flow_notification_queue_capacity = 1024 + ## Jaeger protocol options. [jaeger] ## Whether to enable Jaeger protocol in HTTP API. @@ -250,10 +280,15 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true -## Experimental: enable cumulative OTLP exponential histogram ingestion. -experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Estimated memory budget for cached trace service/operation keys per frontend, +## shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. +## Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. +## Uses 32MiB if memory detection is unavailable. +## Excludes cache and allocator overhead; shared table names are charged per entry. +## @toml2docs:none-default="Auto" +#+ trace_aux_cache_size = "32MiB" ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric ## resource attributes, so metrics-only services reach the semantic graph. experimental_enable_resource_info = false @@ -270,28 +305,14 @@ with_metric_engine = true ## - lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD). ## - unchecked: do not valid strings. prom_validation_mode = "strict" -## Experimental: enable Prometheus remote write v2 native histogram ingestion. -experimental_enable_prometheus_native_histogram = false -## Interval to flush pending rows batcher. -## Set to "0s" to disable batching mode in Prometheus Remote Write endpoint -#+pending_rows_flush_interval = "0s" -## Max rows per pending batch before triggering a flush. -#+max_batch_rows = 100000 -## Max number of concurrent batch flushes. -#+max_concurrent_flushes = 256 -## Capacity of the pending batch worker channel. -#+worker_channel_capacity = 65526 -## Max inflight write requests before backpressure. -#+max_inflight_requests = 3000 -## Maximum number of logical-table flow notifications waiting in the shared queue. -#+flow_notification_queue_capacity = 1024 + ## The WAL options. [wal] ## The provider of the WAL. ## - `raft_engine`: the wal is stored in the local file system by raft-engine. ## - `kafka`: it's remote wal that data is stored in Kafka. -## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** +## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** provider = "raft_engine" ## The directory to store the WAL files. @@ -331,7 +352,8 @@ prefill_log_files = false ## **It's only used when the provider is `raft_engine`**. sync_period = "5s" -## Parallelism during WAL recovery. +## Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. +## @toml2docs:none-default="Auto" recovery_parallelism = 2 ## The Kafka broker endpoints. @@ -406,10 +428,25 @@ overwrite_entry_start_id = false ## The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`. ## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters. -## When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node. +## When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission. ## **It's only used when the provider is `experimental_object_store`**. #+ flush_interval = "100ms" +## When an append to the object store WAL returns. +## - `durable`: an append returns after the object holding its entries is durable (the default). +## - `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog. +## **It's only used when the provider is `experimental_object_store`**. +#+ ack_mode = "durable" + +## The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode. +## Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_bytes = "64MB" + +## The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_age = "8s" + ## What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry. ## - `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default). ## - `fail`: the read fails, so the region does not open. @@ -454,12 +491,6 @@ retry_delay = "500ms" ## If the number of running procedures exceeds this limit, the procedure will be rejected. max_running_procedures = 128 -## flow engine options. -[flow] -## The number of flow worker in flownode. -## Not setting(or set to 0) this value will use the number of CPU cores divided by 2. -#+num_workers=0 - # Example of using S3 as the storage. # [storage] # type = "S3" @@ -501,6 +532,13 @@ max_running_procedures = 128 # credential = "base64-credential" # endpoint = "https://storage.googleapis.com" +# Example of using HDFS as the storage with the native Rust client. +# [storage] +# type = "Hdfs" +# root = "/greptimedb" +# name_node = "hdfs://127.0.0.1:9000" +# options = { "dfs.client.block.write.replace-datanode-on-failure.enable" = "true" } + ## The query engine options. [query] ## Parallelism of the query engine. @@ -509,9 +547,10 @@ parallelism = 0 ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans. +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: @@ -558,6 +597,7 @@ data_home = "./greptimedb_data" ## - `Gcs`: the data is stored in the Google Cloud Storage. ## - `Azblob`: the data is stored in the Azure Blob Storage. ## - `Oss`: the data is stored in the Aliyun OSS. +## - `Hdfs`: the data is stored in the Hadoop Distributed File System. type = "File" ## The S3 bucket name. @@ -565,11 +605,21 @@ type = "File" ## @toml2docs:none-default bucket = "greptimedb" -## The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`. -## **It's only used when the storage type is `S3`, `Oss` and `Azblob`**. +## The directory or object prefix under which data is stored. +## **It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. ## @toml2docs:none-default root = "greptimedb" +## The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +name_node = "hdfs://127.0.0.1:9000" + +## Additional options passed to the native HDFS client. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +options = {} + ## The access key id of the aws account. ## It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key. ## **It's only used when the storage type is `S3` and `Oss`**. @@ -680,7 +730,8 @@ skip_ssl_validation = false ## The Mito engine options. [region_engine.mito] -## Number of region workers. +## Number of region workers. Defaults to half of the available CPU cores, rounded up. +## @toml2docs:none-default="Auto" #+ num_workers = 8 ## Request channel size of each worker. @@ -699,6 +750,13 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Approximate series and range index size limit in open regions, shared across workers. +## Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation +## can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots +## retained by readers are not counted. +## Minimum: 1KiB. Takes effect on restart. +#+ experimental_series_index_max_size = "5GiB" + ## Whether to build and query range indexes when series indexes are enabled. ## Obsolete range-index metadata and files are still cleaned up when disabled. #+ experimental_enable_range_index = false @@ -844,7 +902,7 @@ schedule_compaction_after_edit = true default_flat_format = true ## Whether to enable the experimental two-phase mode for eligible metric series scans. -experimental_series_scan_v2 = true +experimental_series_scan_v2 = false ## The options for index in Mito engine. [region_engine.mito.index] @@ -1015,19 +1073,21 @@ default_ratio = 1.0 ## The slow query log options. [slow_query] ## Whether to enable slow query log. -#+ enable = false +enable = true ## The record type of slow queries. It can be `system_table` or `log`. -## @toml2docs:none-default -#+ record_type = "system_table" +## If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`. +## If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. +record_type = "system_table" -## The threshold of slow query. -## @toml2docs:none-default -#+ threshold = "10s" +## The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. +threshold = "30s" -## The sampling ratio of slow query log. The value should be in the range of (0, 1]. -## @toml2docs:none-default -#+ sample_ratio = 1.0 +## The sampling ratio of slow query log. The value should be in the range of (0, 1]. For example, `0.1` means 10% of the slow queries will be logged and `1.0` means all slow queries will be logged. +sample_ratio = 1.0 + +## The TTL of the `slow_queries` system table. Default is `90d` when `record_type` is `system_table`. +ttl = "90d" ## The tracing options. Only effect when compiled with `tokio-console` feature. #+ [tracing] @@ -1052,6 +1112,7 @@ ttl = "90d" ## `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. +## @toml2docs:none-default="All" #+ event_types = ["create_database"] ## The memory options. diff --git a/cyborg/bin/follow-up-docs-issue.ts b/cyborg/bin/follow-up-docs-issue.ts index acca21fe935..09ef5db2445 100644 --- a/cyborg/bin/follow-up-docs-issue.ts +++ b/cyborg/bin/follow-up-docs-issue.ts @@ -17,7 +17,12 @@ import * as core from '@actions/core' import {handleError, obtainClient} from "@/common"; import {context} from "@actions/github"; -import {PullRequestEditedEvent, PullRequestEvent, PullRequestOpenedEvent} from "@octokit/webhooks-types"; +import { + PullRequestClosedEvent, + PullRequestEditedEvent, + PullRequestEvent, + PullRequestOpenedEvent +} from "@octokit/webhooks-types"; // @ts-expect-error moduleResolution:nodenext issue 54523 import {RequestError} from "@octokit/request-error"; @@ -30,18 +35,59 @@ async function main() { throw new Error(`Only pull request event supported. ${context.eventName} is unsupported.`) } - const client = obtainClient("GITHUB_TOKEN") - const docsClient = obtainClient("DOCS_REPO_TOKEN") const payload = context.payload as PullRequestEvent - const { owner, repo, number, actor, title, html_url } = { + const { owner, repo, number, actor } = { owner: payload.pull_request.base.user.login, repo: payload.pull_request.base.repo.name, number: payload.pull_request.number, - title: payload.pull_request.title, - html_url: payload.pull_request.html_url, actor: payload.pull_request.user.login, } - const followUpDocs = checkPullRequestEvent(payload) + + switch (payload.action) { + case "opened": { + const client = obtainClient("GITHUB_TOKEN") + await updateDocsLabels(client, owner, repo, number, checkPullRequestOpenedEvent(payload as PullRequestOpenedEvent)) + break + } + case "edited": { + const followUpDocs = checkPullRequestEditedEvent(payload as PullRequestEditedEvent) + if (followUpDocs === undefined) { + core.info("Docs checkbox state unchanged; leaving labels as-is.") + return + } + const client = obtainClient("GITHUB_TOKEN") + await updateDocsLabels(client, owner, repo, number, followUpDocs) + break + } + case "closed": { + const event = payload as PullRequestClosedEvent + if (!event.pull_request.merged) { + core.info("PR closed without merging; no docs issue needed.") + return + } + // The docs-required label is the single source of truth. Coordinate + // with pending checkbox-driven label updates before deciding: an edit + // run may still be applying or removing the label, and the event + // payload's label snapshot may be stale. + const client = obtainClient("GITHUB_TOKEN") + const { data: pr } = await client.rest.pulls.get({ owner, repo, pull_number: number }) + await waitForPendingLabelUpdates(client, owner, repo, pr.head.sha) + const { data: freshPr } = await client.rest.pulls.get({ owner, repo, pull_number: number }) + const hasDocsLabel = freshPr.labels.some((label) => label.name === labelDocsRequired) + if (!hasDocsLabel) { + core.info(`Label ${labelDocsRequired} not present; no docs issue needed.`) + return + } + const docsClient = obtainClient("DOCS_REPO_TOKEN") + await createDocsIssue(docsClient, freshPr.title, freshPr.html_url, actor) + break + } + default: + throw new Error(`${payload.action} is unsupported.`) + } +} + +async function updateDocsLabels(client: ReturnType, owner: string, repo: string, number: number, followUpDocs: boolean) { if (followUpDocs) { core.info("Follow up docs.") await client.rest.issues.removeLabel({ @@ -55,28 +101,6 @@ async function main() { await client.rest.issues.addLabels({ owner, repo, issue_number: number, labels: [labelDocsRequired], }) - - // Get available assignees for the docs repo - const assigneesResponse = await docsClient.rest.issues.listAssignees({ - owner: 'GreptimeTeam', - repo: 'docs', - }) - const validAssignees = assigneesResponse.data.map(assignee => assignee.login) - core.info(`Available assignees: ${validAssignees.join(', ')}`) - - // Check if the actor is a valid assignee, otherwise fallback to fengjiachun - const assignee = validAssignees.includes(actor) ? actor : 'fengjiachun' - core.info(`Assigning issue to: ${assignee}`) - - await docsClient.rest.issues.create({ - owner: 'GreptimeTeam', - repo: 'docs', - title: `Update docs for ${title}`, - body: `A document change request is generated from ${html_url}`, - assignee: assignee, - }).then((res) => { - core.info(`Created issue ${res.data}`) - }) } else { core.info("No need to follow up docs.") await client.rest.issues.removeLabel({ @@ -93,15 +117,52 @@ async function main() { } } -function checkPullRequestEvent(payload: PullRequestEvent) { - switch (payload.action) { - case "opened": - return checkPullRequestOpenedEvent(payload as PullRequestOpenedEvent) - case "edited": - return checkPullRequestEditedEvent(payload as PullRequestEditedEvent) - default: - throw new Error(`${payload.action} is unsupported.`) +async function createDocsIssue(docsClient: ReturnType, title: string, html_url: string, actor: string) { + core.info("Creating follow-up docs issue for merged PR.") + + // Get available assignees for the docs repo + const assigneesResponse = await docsClient.rest.issues.listAssignees({ + owner: 'GreptimeTeam', + repo: 'docs', + }) + const validAssignees = assigneesResponse.data.map(assignee => assignee.login) + core.info(`Available assignees: ${validAssignees.join(', ')}`) + + // Check if the actor is a valid assignee, otherwise fallback to fengjiachun + const assignee = validAssignees.includes(actor) ? actor : 'fengjiachun' + core.info(`Assigning issue to: ${assignee}`) + + await docsClient.rest.issues.create({ + owner: 'GreptimeTeam', + repo: 'docs', + title: `Update docs for ${title}`, + body: `A document change request is generated from ${html_url}`, + assignee: assignee, + }).then((res) => { + core.info(`Created issue ${res.data}`) + }) +} + +// Waits until no other runs of this workflow are pending for the PR head SHA, +// so that any checkbox-driven label update has been applied before the caller +// reads the labels. Bounded by a deadline; on timeout it proceeds with whatever +// the current label state is. +async function waitForPendingLabelUpdates(client: ReturnType, owner: string, repo: string, headSha: string) { + const deadlineMs = 5 * 60 * 1000 + const pollIntervalMs = 5000 + const deadline = Date.now() + deadlineMs + while (Date.now() < deadline) { + const { data } = await client.rest.actions.listWorkflowRuns({ + owner, repo, workflow_id: "docbot.yml", head_sha: headSha, per_page: 20, + }) + const pending = data.workflow_runs.filter((run) => run.id !== context.runId && run.status !== "completed") + if (pending.length === 0) { + return + } + core.info(`Waiting for ${pending.length} pending docbot run(s) to finish label updates...`) + await new Promise((resolve) => setTimeout(resolve, pollIntervalMs)) } + core.warning("Timed out waiting for pending docbot runs; proceeding with the current label state.") } function checkPullRequestOpenedEvent(event: PullRequestOpenedEvent): boolean { @@ -109,11 +170,19 @@ function checkPullRequestOpenedEvent(event: PullRequestOpenedEvent): boolean { return event.pull_request.body?.includes(needFollowUpDocs) } -function checkPullRequestEditedEvent(event: PullRequestEditedEvent): boolean { - const previous = event.changes.body?.from.includes(needFollowUpDocs) - const current = event.pull_request.body?.includes(needFollowUpDocs) - // from docs-not-need to docs-required - return (!previous) && current +// Returns undefined when the checkbox state did not change in this edit, so the +// caller leaves the labels untouched (preserving manual label overrides). +function checkPullRequestEditedEvent(event: PullRequestEditedEvent): boolean | undefined { + if (!event.changes.body) { + // The body was not part of this edit (e.g. title-only edit). + return undefined + } + const previous = event.changes.body.from.includes(needFollowUpDocs) + const current = event.pull_request.body?.includes(needFollowUpDocs) ?? false + if (previous === current) { + return undefined + } + return current } main().catch(handleError) diff --git a/docs/how-to/ai-functions.md b/docs/how-to/ai-functions.md new file mode 100644 index 00000000000..5c33ffff50e --- /dev/null +++ b/docs/how-to/ai-functions.md @@ -0,0 +1,247 @@ +# AI functions (experimental MVP) + +Three asynchronous SQL scalar functions provide natural-language matching, +classification, and rating. The current backend is TypeSafe's Jev model: + +| Function | Mode | Result | +| --- | --- | --- | +| `ai_match(text, prompt)` | Noul | `Float64` probability in **[0, 1]** that the prompt's statement is true | +| `ai_choose(text, prompt, criteria)` | Choice | `String` containing the selected option's name | +| `ai_score(text, prompt, criteria)` | Score | JSONB object containing `score`, `confidence`, and a `probabilities` array | + +Use scalar results or extracted JSON fields in `SELECT`, comparisons in `WHERE`, +and `ORDER BY`. +All arguments are SQL strings. Any SQL NULL argument produces SQL `NULL` +without validating that row's criteria or calling the API. + +## Start GreptimeDB + +The functions are compiled and registered by the **default-enabled Cargo feature +`ai_functions`**. Enable the Jev backend separately in the server process: + +```sh +export GREPTIMEDB_EXPERIMENTAL_JEV=true +export JEV_API_KEY='' +cargo run -p cmd -- standalone start +``` + +If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first. +The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to +`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL). +The MVP uses environment variables rather than TOML configuration. +The Cargo feature includes all three SQL functions in the build; the runtime environment +variables enable and configure its API calls. Default compilation does not turn +on external API calls. To enable the feature explicitly, use +`--features ai_functions`. + +## Query + +For a new example table, quote `service` and `message` in the DDL because the +current parser treats them as keywords: + +```sql +CREATE TABLE events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); +``` + +### Matching probability (Noul) + +```sql +SELECT occurred_at, service, message +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match( + (message), + 'The event reports that a payment still failed after retries.' + ) >= 0.8 +ORDER BY occurred_at; +``` + +To return the scores and rank matching events: + +```sql +SELECT occurred_at, message, + ai_match(message, 'The event reports that a payment still failed after retries.') AS score +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' +ORDER BY score DESC NULLS LAST; +``` + +Both arguments are strings; larger scores indicate a stronger match. +`(message)` is an ordinary parenthesized string expression. For several columns, +combine them explicitly, for example `concat(service, ': ', message)`. +Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call. + +### Classification (Choice) + +`criteria` is a JSON object with **1 to 255 options**. Each key is an option name; +its value is a description (string, object, or array), or JSON `null` if the name +is sufficient. The function returns the selected key, suitable for filtering +or grouping: + +```sql +SELECT occurred_at, message, + ai_choose(message, 'Which team should handle this event?', + '{"billing":"Payments, invoices, refunds","technical":"Bugs and outages","other":null}') AS team +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z'; +``` + +### Rating (Score) + +`criteria` is a JSON array of **2 to 10 ordered level descriptions**, from low +to high. Descriptions can be strings, objects, or arrays. Level numbers start +at zero. The function returns a JSONB object from one model evaluation: + +```json +{ + "score": 1.05, + "confidence": 0.92, + "probabilities": [0.0, 0.95, 0.05] +} +``` + +- `score` is the provider's probability-weighted mean of the level numbers, in + `[0, N − 1]` for N levels. It may fall between levels. +- `confidence` is the provider's confidence in `[0, 1]`, reflecting how concentrated + the distribution is. It is not a guarantee that the answer is correct. +- `probabilities[i]` is the probability of `criteria[i]`, with exactly N entries in + level order. These are individual level probabilities, not cumulative probabilities. + +Different distributions can have the same score. `[0, 1, 0]` and `[0.5, 0, 0.5]` +both have score `1`: the first concentrates on the middle level, while the second +is split between the extremes. Inspect confidence and the distribution before +treating a mean score as a definite severity level. + +For example, compute the rating once, extract numeric fields, and rank only rows +meeting a chosen confidence threshold: + +```sql +SELECT occurred_at, message, + json_get_float(rating, 'score') AS severity, + json_get_float(rating, 'confidence') AS confidence +FROM ( + SELECT occurred_at, message, + ai_score(message, 'How severe is this event?', + '["No impact to functionality","Degraded service with a workaround","Blocking issue with no workaround"]') AS rating + FROM events + WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' +) AS rated +WHERE json_get_float(rating, 'confidence') >= 0.8 +ORDER BY severity DESC NULLS LAST; +``` + +Use `json_to_string(rating)` to display the full object, or +`json_get_float(rating, 'probabilities[2]')` to read the probability of level 2. +The array positions correspond to the supplied criteria, which provide the level +descriptions. + +## Reusing an evaluation + +AI functions are marked volatile to prevent HTTP calls during planning. Identical +calls written separately in `SELECT` and `WHERE` are evaluated separately. If the +filter evaluates N non-null rows and M rows pass to the projection, this can cost +N + M requests, reaching 2N when all rows pass. + +To return a result and filter by it, compute it once in a subquery and reference +its alias in the outer query: + +```sql +SELECT occurred_at, message, score +FROM ( + SELECT occurred_at, message, + ai_match(message, 'The event reports that a payment still failed after retries.') AS score + FROM events + WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' +) AS scored +WHERE score >= 0.8 +ORDER BY score DESC; +``` + +This also applies to `ai_choose` and `ai_score`. Extract multiple JSON fields from +the aliased `ai_score` result, as in the rating example above, rather than calling +the model again for each field. Filter pushdown preserves the volatile projection +instead of duplicating it into the outer predicate. + +## MVP behavior + +- Each non-null row makes one HTTP request per function invocation: the text is + `state`, and the prompt is the question's `instructions`. Choice and Score pass + the parsed JSON as `criteria`. Calls to different functions are separate requests. +- Criteria for all non-null rows in a batch are validated before sending that + batch's requests. Malformed JSON, invalid description types, or invalid option/ + level counts fail the query locally. +- Answers must have the requested question type. Noul probabilities outside + `[0, 1]`, Choice labels not in the criteria, and Score values outside `[0, N − 1]` + fail the query rather than being clamped or replaced with NULL. Score responses + must also contain a numeric confidence in `[0, 1]` and probabilities for every + level. Each probability must be numeric and in `[0, 1]`; their sum must be within + `1e-6` of 1. Provider values are preserved without renormalizing the distribution + or recomputing the score. +- Up to eight requests run concurrently per expression/batch invocation, not per + query or process. Concurrent partitions and queries can exceed eight requests + in total. Each request has a 30-second timeout. Errors (including rate limits + and invalid responses) fail the query. There is no automatic retry or + cross-query cache. +- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the + candidate set, but SQL does not guarantee left-to-right predicate evaluation or + that `LIMIT` bounds the number of API calls. +- Text passed to the function is sent to the configured TypeSafe endpoint. + The API key remains in the server environment, not in SQL. +- This MVP is intended for standalone use. A distributed deployment would need + the environment on every node evaluating the function and separate validation. + +API reference: + +## Before stabilizing + +The experimental MVP still needs the following controls before stabilization: + +- A process-wide concurrency limit shared by AI function invocations, such as a semaphore. +- Bounded retries with exponential backoff for HTTP `429` and `529`, following + TypeSafe's rate-limit guidance. +- Request budgets and metrics for API calls, latency, retries, and rate-limit errors. + +These are follow-up work, not guarantees provided by the current implementation. + +## Validation + +Check both the default build (where all three functions are registered) and an isolated +`common-function` build without default features. The regular AI function tests use a local +HTTP server and need no API key: + +```sh +cargo nextest run -p common-function +cargo nextest run -p common-function --no-default-features +cargo nextest run -p common-function --no-default-features --features ai_functions +``` + +The sqlness runner explicitly includes `ai_functions` when building its test binary: + +```sh +cargo sqlness bare -t ai_functions +``` + +If using `--bins-dir`, provide a binary built with `ai_functions` (included by +default). CI covers the AI-enabled path through the regular unit tests and +sqlness tests; the feature-off checks above can be run locally when needed. + +An opt-in test calls the real service on synthetic payment events: + +```sh +source ~/.zshrc +GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \ + -E 'test(test_ai_match_live)' --run-ignored ignored-only +``` diff --git a/docs/how-to/jev.md b/docs/how-to/jev.md deleted file mode 100644 index 2a21f085946..00000000000 --- a/docs/how-to/jev.md +++ /dev/null @@ -1,119 +0,0 @@ -# Natural-language filtering with Jev (experimental MVP) - -`jev(text, statement, threshold)` returns whether Jev judges the statement true -for the text with probability **greater than or equal to** the threshold. -It is an asynchronous SQL scalar function usable in `WHERE` and `SELECT`. - -## Start GreptimeDB - -Jev is compiled and registered by the **default-enabled Cargo feature -`ai_functions`**. Enable API evaluation separately in the server process: - -```sh -export GREPTIMEDB_EXPERIMENTAL_JEV=true -export JEV_API_KEY='' -cargo run -p cmd -- standalone start -``` - -If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first. -The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to -`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL). -The MVP uses environment variables rather than TOML configuration. -The Cargo feature includes the SQL function in the build; the runtime environment -variables enable and configure its API calls. Default compilation does not turn -on external API calls. To enable the feature explicitly, use -`--features ai_functions`. - -## Query - -For a new example table, quote `service` and `message` in the DDL because the -current parser treats them as keywords: - -```sql -CREATE TABLE events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); -``` - -```sql -SELECT occurred_at, service, message -FROM events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev( - (message), - 'The event reports that a payment still failed after retries.', - 0.8 - ) -ORDER BY occurred_at; -``` - -The first two arguments are strings and the third is a number in `[0, 1]`. -`(message)` is an ordinary parenthesized string expression. For several columns, -combine them explicitly, for example `concat(service, ': ', message)`. -Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call. - -## MVP behavior - -- Each non-null row makes one HTTP request: the text is `state`, and the statement - is a `noul` question's `instructions`. The returned `noul` probability is compared - with the threshold locally; this is not the separate Choice/Score confidence. -- Up to eight requests run concurrently per expression/batch invocation, not per - query or process. Concurrent partitions and queries can exceed eight requests - in total. Each request has a 30-second timeout. Errors (including rate limits - and invalid responses) fail the query. There is no automatic retry or - cross-query cache. -- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the - candidate set, but SQL does not guarantee left-to-right predicate evaluation or - that `LIMIT` bounds the number of API calls. -- Text passed to the function is sent to the configured TypeSafe endpoint. - The API key remains in the server environment, not in SQL. -- This MVP is intended for standalone use. A distributed deployment would need - the environment on every node evaluating the function and separate validation. - -API reference: - -## Before stabilizing - -The experimental MVP still needs the following controls before stabilization: - -- A process-wide concurrency limit shared by Jev invocations, such as a semaphore. -- Bounded retries with exponential backoff for HTTP `429` and `529`, following - TypeSafe's rate-limit guidance. -- Request budgets and metrics for API calls, latency, retries, and rate-limit errors. - -These are follow-up work, not guarantees provided by the current implementation. - -## Validation - -Check both the default build (where `jev` is registered) and an isolated -`common-function` build without default features. Jev's regular tests use a local -HTTP server and need no API key: - -```sh -cargo nextest run -p common-function -cargo nextest run -p common-function --no-default-features -cargo nextest run -p common-function --no-default-features --features ai_functions -``` - -The sqlness runner explicitly includes `ai_functions` when building its test binary: - -```sh -cargo sqlness bare -t jev -``` - -If using `--bins-dir`, provide a binary built with `ai_functions` (included by -default). CI covers the AI-enabled path through the regular unit tests and -sqlness tests; the feature-off checks above can be run locally when needed. - -An opt-in test calls the real service on synthetic payment events: - -```sh -source ~/.zshrc -GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \ - -E 'test(test_jev_live)' --run-ignored ignored-only -``` diff --git a/docs/rfcs/2026-05-28-table-semantic-layer.md b/docs/rfcs/2026-05-28-table-semantic-layer.md index 24ed96dc81b..cb5f6e84dc4 100644 --- a/docs/rfcs/2026-05-28-table-semantic-layer.md +++ b/docs/rfcs/2026-05-28-table-semantic-layer.md @@ -93,13 +93,25 @@ Two design decisions worth pinning down up front, because they constrain everyth - **Conflict.** Some table-level keys (`trace.conventions` lifted from `schema_url`, `metric.temporality`, ...) cannot represent the truth when a long-lived table sees rows from multiple sources. v1 records `mixed` or `unknown` rather than a fictitious single value. Downstream consumers must treat any single-valued semantic key as best-effort, not strong evidence. - **Update.** Semantic options are stamped at table creation. v1 does not specify an update path; promoting `metadata_quality` from `inferred` to `declared`, refreshing `resource.attributes_preserved`, or revising `trace.conventions` on later writes is deferred. If real usage shows update is needed, it lands as a separate RFC. +OTLP delta sums and explicit histograms store raw interval values without +accumulating across timestamps. Delta-to-cumulative conversion and its +per-series ingestion state are intentionally out of scope. Explicit histogram +bucket counts are prefix-summed only within each point to produce the classic +histogram representation. Delta exponential histograms are currently rejected; +[planned raw-delta support](2026-08-04-native-histograms.md#planned-raw-delta-exponential-histograms) +will store their interval values in the canonical native-histogram Struct using +the same temporality tag and a gauge reset hint. The tag will select histogram +summation for `increase()` and summation divided by range duration for `rate()`, +without cumulative reset correction or extrapolation. + OTLP delta sums and explicit histograms additionally store the query-visible String tag `otlp_aggregation_temporality="delta"` on each generated row. Its name is fixed and does not follow `default_column_prefix`. The tag is part of series identity and is authoritative for per-series float `rate()` and `increase()` behavior; the table option is never used as a row-level -discriminator. Native histograms retain their native algorithms. Prometheus -metadata reports `unknown` for counter, histogram, and up/down-counter tables +discriminator. Native histograms currently retain their native algorithms; +the planned delta extension will change only tagged delta histogram series. +Prometheus metadata reports `unknown` for counter, histogram, and up/down-counter tables whose catalog temporality is `delta` or `mixed`. A same-request conflict can create `mixed`; a later write does not update an existing table option, so the catalog value can also be stale while the row tag remains authoritative. diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md index c714b35351b..3aef4eb7ffb 100644 --- a/docs/rfcs/2026-08-04-native-histograms.md +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -1,6 +1,6 @@ --- Feature Name: Native Histogram Support and Compatibility Decisions -Tracking Issue: TBD +Tracking Issue: https://github.com/GreptimeTeam/greptimedb/issues/8887 Date: 2026-08-04 Author: codex --- @@ -13,8 +13,8 @@ field and evaluates them as first-class PromQL samples. Prometheus Remote Write ingestion transport: accepted cumulative points are normalized into the same Struct before persistence and are queried only as native histograms. -Native histograms are experimental. Prometheus Remote Write 2.0 and cumulative -OTLP exponential histograms are supported behind separate configuration gates. +Prometheus Remote Write 2.0 native histograms and cumulative OTLP/HTTP +exponential histograms are enabled by default. Remote Write 1.0 histogram payloads are rejected instead of being acknowledged and dropped. Native-histogram Remote Read is deferred; the existing Remote Read path continues to return scalar samples only. @@ -34,6 +34,9 @@ continues to return scalar samples only. - Returning native histograms through Prometheus Remote Read. - Persisting Remote Write metric metadata. - Persisting exemplars. +- Converting OTLP delta values to cumulative values during ingestion. Supported + delta metrics are stored as interval values without per-series accumulation + state; this also applies to the planned exponential-histogram support. - Providing an OTLP-specific histogram query surface or reconstructing and re-exporting the original OTLP point after persistence. - Removing mixed float/histogram handling from PromQL expressions. @@ -95,8 +98,7 @@ that send native histograms must use Remote Write 2.0. ## Remote Write 2.0 -Remote Write 2.0 accepts integer and float native histograms while -`prom_store.experimental_enable_prometheus_native_histogram` is enabled. Supported +Remote Write 2.0 accepts integer and float native histograms by default. Supported exponential schemas are `-4` through `8`; schema `-53` represents native histograms with custom buckets. @@ -121,15 +123,23 @@ explicit if sampled responses are implemented. ## OTLP -OTLP exponential histograms are accepted when -`otlp.experimental_enable_exponential_histogram` is enabled. The option defaults -to false and applies to OTLP/HTTP. Disabled points are rejected rather than -silently acknowledged. OTel Arrow exponential histograms are rejected because +OTLP/HTTP exponential histograms are accepted by default. +OTel Arrow exponential histograms are rejected because the current Arrow wire format omits `zero_threshold`; accepting them would -silently change the distribution. Cumulative temporality is required; delta and -unspecified exponential histograms are rejected before their points are +silently change the distribution. Currently, cumulative temporality is required; +delta and unspecified exponential histograms are rejected before their points are converted. Explicit OTLP histograms keep their existing `_bucket`, `_sum`, and -`_count` representation, including their existing delta behavior. +`_count` representation. OTLP delta sums and explicit histograms are stored as +raw interval values with `otlp_aggregation_temporality="delta"`. Explicit +histogram bucket counts are prefix-summed within each point for the classic +histogram representation, but values are never accumulated across timestamps. +For these tagged float series, PromQL `increase()` sums the interval values +in the query range, and `rate()` divides that sum by the range duration. +Delta-to-cumulative conversion is intentionally out of scope, not deferred work. +Raw-delta exponential-histogram support is planned below; delta exponential +points remain rejected until it is implemented. See +[Table Semantic Layer](2026-05-28-table-semantic-layer.md#conflict-and-update-semantics) +for the per-series temporality contract. OTLP scales `-4` through `8` map directly to Prometheus schemas. Higher scales are downscaled to schema `8`: dense counts that collide are merged before the @@ -159,8 +169,36 @@ for mixed batches and `INVALID_ARGUMENT` when all points are rejected. Rejection details are bounded, and metric metadata is emitted only for a metric that produced an accepted row. -Minimum, maximum, and exemplars are not persisted. Delta accumulation, -zero-run span compaction, and a dedicated rejection metric remain deferred. +Zero-count buckets are omitted after downscaling, splitting nonzero runs into +sparse spans without changing integer counts. Rejected data points increment +`greptime_servers_otlp_exponential_histogram_rejected_data_points_total`, labeled +by `reason`: `delta_temporality`, `unspecified_temporality`, or +`invalid_data_point`. Each rejected point is counted once, including when +resource descriptors are enabled. Minimum, maximum, and exemplars are not persisted. + +### Planned raw-delta exponential histograms + +Accept OTLP delta exponential histograms and store each point's interval counts +and sum in the existing canonical native-histogram Struct, without accumulating +across timestamps. Reuse the integer-count encoding, scale conversion, +zero-threshold preservation, stale-marker handling, and partial-success path. +OTel Arrow support remains blocked independently by its missing `zero_threshold`. + +Store `otlp_aggregation_temporality="delta"` as part of series identity, as for +delta sums and explicit histograms. Use the gauge reset hint for delta histogram +samples so decreasing interval counts are not interpreted as counter resets. +Prometheus uses the same gauge hint for +[raw-delta exponential-histogram ingestion](https://github.com/prometheus/prometheus/blob/9f159f6df198280b757163ec555ffde0c8bfe8e0/storage/remote/otlptranslator/prometheusremotewrite/histograms.go#L125). +The temporality tag, rather than the gauge hint alone or table metadata, +determines raw-delta query behavior. + +Extend the per-series `increase()` and `rate()` selection to native histogram +inputs: sum interval histograms over the query range, then divide by the range +duration for `rate()`. Do not apply cumulative reset correction, synthetic-zero +insertion, or counter extrapolation to these delta samples. Preserve cumulative +histogram behavior and mixed float/histogram range warnings. Acceptance coverage +must include multiple intervals whose counts decrease, exact integer ingestion, +stale points, partial rejection, and pure-native and mixed-range query paths. # PromQL Compatibility Decisions @@ -276,7 +314,8 @@ native-histogram data remains readable. - Native-histogram Remote Read, including exact integer round-trips and streamed chunks with start timestamps. -- OTLP exponential-histogram delta accumulation. +- Raw-delta OTLP exponential-histogram storage and temporality-aware native + histogram queries, as described above; no delta-to-cumulative conversion. - Persistent Remote Write metadata and accurate help/unit updates. - Native-histogram exemplars and exemplar query APIs. - Start-timestamp overlap annotations. @@ -284,3 +323,10 @@ native-histogram data remains readable. - Versioned two-phase native-histogram aggregation state. - Exact Prometheus summation compensation if measured precision differences justify the extra aggregate state. + +## Configuration migration + +The `prom_store.experimental_enable_prometheus_native_histogram` and +`otlp.experimental_enable_exponential_histogram` options have been removed. +Remove them from existing configurations; both ingestion paths are always enabled +when their protocol is enabled. Old option values are ignored, including `false`. diff --git a/docs/rfcs/2026-09-11-metric-export-import.md b/docs/rfcs/2026-09-11-metric-export-import.md index 3d2c93e4695..8ddc31f40b0 100644 --- a/docs/rfcs/2026-09-11-metric-export-import.md +++ b/docs/rfcs/2026-09-11-metric-export-import.md @@ -390,14 +390,34 @@ validation, schema-option inheritance and target metadata resolution. Bound the request size and the importer's DDL buffering separately. The production CLI needs an authenticated server operation that checks CREATE -authorization for every member and reuses the ordinary DDL path. Transport, -capability discovery and request limits must be agreed before accepting this -interface. Initially bound a batch to 128 statements and 1 MiB of SQL, with an -explicit oversized-statement error; separately bound the importer's parsed DDL -buffer. A target without batch-DDL capability uses ordinary DDL if it supports +authorization for every member and reuses the ordinary DDL path. Bound the +importer's parsed DDL buffer separately from the server request limits below. +A target without batch-DDL capability uses ordinary DDL if it supports the snapshot's data layout. Do not retry an ambiguously completed batch through another path. +The server operation is `POST /v1/ddl/logical-tables` with an +`application/x-www-form-urlencoded` body containing `sql` and optional `db`. +`db` uses the same catalog/schema interpretation and defaults as `/v1/sql`. +Authenticated `GET /v1/capabilities` advertises `metric_batch_ddl: 1` when the +SQL handler supports the operation, independently of packed import/export. +The form body is limited to 4 MiB before decoding, or the configured HTTP body +limit when that limit is stricter. Decoded SQL is limited to 1 MiB and 1–128 +statements. Only explicit Metric logical CREATE TABLE statements in one +catalog/schema with one literal `on_physical_table` value are accepted; +duplicate normalized target names are rejected. + +All members pass SQL interception, authorization and ordinary CREATE preparation +before one logical-table batch procedure is submitted. Success uses the normal +GreptimeDB v1 response with one `affectedrows: 0` output per input statement in +order; this is CREATE's output, not a count of newly created tables. A failure +returns one error without a partial success list. Validation failures before +submission have no DDL side effects. Once submitted, the existing procedure may +have applied some or all changes even if the HTTP request times out or disconnects; +there is no cross-table rollback guarantee. `IF NOT EXISTS` keeps ordinary CREATE +semantics. The operation does not regroup requests or retry through single-table +CREATE. + ## Packed data COPY The CLI retains one `COPY DATABASE FROM` operation per `(chunk, schema)` task. diff --git a/flake.lock b/flake.lock index 27c1797f10b..621f1e2f3a0 100644 --- a/flake.lock +++ b/flake.lock @@ -8,11 +8,11 @@ "rust-analyzer-src": "rust-analyzer-src" }, "locked": { - "lastModified": 1785314864, - "narHash": "sha256-imz9J5iMNersDeQ1ymenFg35N6C2R3MCaGg2i0XCIrE=", + "lastModified": 1790642863, + "narHash": "sha256-JIwE6EnebvxW7QFFEitWHuZDXwN6Qcf6ZC8DIvYW3y8=", "owner": "nix-community", "repo": "fenix", - "rev": "594418b2c9ea0731bee6988236ef5bfd97e81dca", + "rev": "a8f386c482e78e9a48be5abecdc17a9b098a95a1", "type": "github" }, "original": { @@ -41,11 +41,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1785133411, - "narHash": "sha256-Yjv0WEg39KRYS0rBdTbu6Fc/or/ihAKk13W9sQ6VWd0=", + "lastModified": 1790533781, + "narHash": "sha256-F+b8u43BVF5DlM0pz7JQSbtbwKy5PrHIT5zQUekXsD0=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "2f5a153c270b70cb0f8c11f46d96d6d3bc39f4e3", + "rev": "cf5e76507c6e23b59f7e0ffcc7baa2a39ddd8442", "type": "github" }, "original": { @@ -65,11 +65,11 @@ "rust-analyzer-src": { "flake": false, "locked": { - "lastModified": 1785261141, - "narHash": "sha256-sA+DHPejWD68mLA7gtTiHGd5T41F6W+NKf0g+ibFCW4=", + "lastModified": 1790532271, + "narHash": "sha256-j8y/AO3AMDF9p+11aIDhnA21JnNj5PjvcejTI43vI8Q=", "owner": "rust-lang", "repo": "rust-analyzer", - "rev": "bec66814323579659ffd77c909b3d963af118ece", + "rev": "03fcb77246f2568adb0e9b2fa60d19c6cc1686f4", "type": "github" }, "original": { diff --git a/flake.nix b/flake.nix index a5886b8add1..df71403b24f 100644 --- a/flake.nix +++ b/flake.nix @@ -22,7 +22,7 @@ lib = nixpkgs.lib; rustToolchain = fenix.packages.${system}.fromToolchainName { name = (lib.importTOML ./rust-toolchain.toml).toolchain.channel; - sha256 = "sha256-rboGKQLH4eDuiY01SINOqmXUFUNr9F4awoFZGzib17o="; + sha256 = "sha256-p8h3Sl/YRByZfZTAKXdsvF6xEenXKrXSVvpphmZENH4="; }; in { @@ -35,6 +35,7 @@ protobuf gnumake mold + wild (rustToolchain.withComponents [ "cargo" "clippy" @@ -46,7 +47,7 @@ ]) cargo-nextest cargo-llvm-cov - cargo-udeps + cargo-shear taplo curl gnuplot ## for cargo bench diff --git a/grafana/dashboards/metrics/cluster/dashboard.json b/grafana/dashboards/metrics/cluster/dashboard.json index bd22a0da597..daa0f018836 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.json +++ b/grafana/dashboards/metrics/cluster/dashboard.json @@ -3585,19 +3585,6 @@ }, "pluginVersion": "11.6.0", "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by (code) (rate(greptime_flow_errors[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "flow-{{code}}", - "range": true, - "refId": "A" - }, { "datasource": { "type": "prometheus", @@ -11583,7 +11570,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "hit-{{type}}", @@ -11596,7 +11583,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "miss-{{type}}", @@ -11609,7 +11596,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "evict-{{type}}-{{cause}}", @@ -16622,127 +16609,6 @@ "title": "Flow Ingest / Output Rate", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Ingest Latency.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "points", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 170 - }, - "id": 521, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p95", - "range": true, - "refId": "A" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p99", - "range": true, - "refId": "B" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-avg", - "range": true, - "refId": "C" - } - ], - "title": "Flow Ingest Latency", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -16805,7 +16671,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, + "x": 8, "y": 170 }, "id": 522, @@ -16851,196 +16717,6 @@ "title": "Flow Operation Latency", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Buffer Size per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 0, - "y": 178 - }, - "id": 523, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "greptime_flow_input_buf_size", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Buffer Size per Instance", - "type": "timeseries" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Processing Error per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 178 - }, - "id": 524, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-[{{code}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Processing Error per Instance", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -17105,7 +16781,7 @@ "h": 8, "w": 8, "x": 16, - "y": 178 + "y": 170 }, "id": 558, "options": { @@ -17229,7 +16905,7 @@ "h": 8, "w": 8, "x": 0, - "y": 186 + "y": 178 }, "id": 559, "options": { @@ -17353,7 +17029,7 @@ "h": 8, "w": 8, "x": 8, - "y": 186 + "y": 178 }, "id": 560, "options": { @@ -17464,7 +17140,7 @@ "h": 8, "w": 8, "x": 16, - "y": 186 + "y": 178 }, "id": 561, "options": { @@ -17575,7 +17251,7 @@ "h": 8, "w": 8, "x": 0, - "y": 194 + "y": 186 }, "id": 562, "options": { @@ -17673,7 +17349,7 @@ "h": 8, "w": 8, "x": 8, - "y": 194 + "y": 186 }, "id": 563, "options": { @@ -17771,7 +17447,7 @@ "h": 8, "w": 8, "x": 16, - "y": 194 + "y": 186 }, "id": 564, "options": { diff --git a/grafana/dashboards/metrics/cluster/dashboard.md b/grafana/dashboards/metrics/cluster/dashboard.md index b8e050114c6..3278db52bfd 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.md +++ b/grafana/dashboards/metrics/cluster/dashboard.md @@ -36,7 +36,7 @@ | Scan and Compaction Memory Rejects | `sum(rate(greptime_mito_scan_requests_rejected_total{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_scan_memory_exhausted_total{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_compaction_memory_rejected_total{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Datanode scan and compaction memory rejection/exhaustion counters. | `prometheus` | `rps` | `scan-rejected` | | OpenDAL Errors | `sum by (scheme, operation, error) (rate(opendal_operation_errors_total{instance=~"$datanode",error!="NotFound"}[$__rate_interval]))` | `timeseries` | Object-store errors by scheme, operation, and error, excluding NotFound noise. | `prometheus` | `eps` | `{{scheme}}-{{operation}}-{{error}}` | | Metasrv Failures | `sum(rate(greptime_meta_region_migration_fail[$__rate_interval]))`
`sum(rate(greptime_meta_reconciliation_procedure_error[$__rate_interval]))` | `timeseries` | Region migration and reconciliation failures in metasrv. | `prometheus` | `eps` | `migration-fail` | -| Flow and Trigger Failures | `sum by (code) (rate(greptime_flow_errors[$__rate_interval]))`
`sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `flow-{{code}}` | +| Flow and Trigger Failures | `sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `trigger-evaluate` | | Mito GC Failures | `sum(rate(greptime_mito_gc_errors_total{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_gc_orphaned_index_files{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_gc_skipped_unparsable_files{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Mito garbage-collection errors and skipped/orphaned files on datanodes. | `prometheus` | `short` | `gc-errors` | # Capacity | Title | Query | Type | Description | Datasource | Unit | Legend Format | @@ -131,7 +131,7 @@ | Index Memory Usage | `greptime_index_apply_memory_usage`
`sum by (type) (greptime_index_create_memory_usage)` | `timeseries` | Memory used while applying and creating indexes. Growth here can explain memory pressure during indexed flush or compaction work. | `prometheus` | `bytes` | `apply` | | Index IO Bytes | `sum by (type, file_type) (rate(greptime_index_io_bytes_total[$__rate_interval]))` | `timeseries` | Index read and write byte rates by operation and file type for puffin and intermediate files. | `prometheus` | `Bps` | `{{type}}-{{file_type}}` | | Index IO Operations | `sum by (type, file_type) (rate(greptime_index_io_op_total[$__rate_interval]))` | `timeseries` | Index IO operation rates by operation and file type, including read, write, seek, and flush operations. | `prometheus` | `ops` | `{{type}}-{{file_type}}` | -| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | +| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | # Metasrv | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | @@ -278,10 +278,7 @@ ORDER BY data_size DESC;` | `piechart` | Distribution of leader regions and data | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | | Flow Ingest / Output Rate | `sum by(instance, pod, direction) (rate(greptime_flow_processed_rows[$__rate_interval]))` | `timeseries` | Flow Ingest / Output Rate. | `prometheus` | -- | `[{{pod}}]-[{{instance}}]-[{{direction}}]` | -| Flow Ingest Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))` | `timeseries` | Flow Ingest Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-p95` | | Flow Operation Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))`
`histogram_quantile(0.99, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))` | `timeseries` | Flow Operation Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{type}}]-p95` | -| Flow Buffer Size per Instance | `greptime_flow_input_buf_size` | `timeseries` | Flow Buffer Size per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]` | -| Flow Processing Error per Instance | `sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))` | `timeseries` | Flow Processing Error per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{code}}]` | | Flow Batching Attempt / Error Rate | `sum by(instance,pod,flow_id) (rate(greptime_flow_batching_start_query_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_error_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_slow_query_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode task attempt, error, and slow-query rates by flow. | `prometheus` | `ops` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-attempt` | | Flow Batching Query Latency | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`histogram_quantile(0.99, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode query latency by flow. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | | Flow Batching Query Window Count | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_count[$__rate_interval]))` | `timeseries` | Flow batching mode query window count by flow. | `prometheus` | `short` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | diff --git a/grafana/dashboards/metrics/cluster/dashboard.yaml b/grafana/dashboards/metrics/cluster/dashboard.yaml index cd105edfcd2..bc659b401b5 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.yaml +++ b/grafana/dashboards/metrics/cluster/dashboard.yaml @@ -515,11 +515,6 @@ groups: description: Derived-data and alerting pipeline failures. unit: eps queries: - - expr: sum by (code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: flow-{{code}} - expr: sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval])) datasource: type: prometheus @@ -1572,17 +1567,17 @@ groups: description: Index-related cache hits, misses, and evictions from Mito caches. unit: ops queries: - - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: hit-{{type}} - - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: miss-{{type}} - - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} @@ -2199,25 +2194,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{pod}}]-[{{instance}}]-[{{direction}}]' - - title: Flow Ingest Latency - type: timeseries - description: Flow Ingest Latency. - queries: - - expr: histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p95' - - expr: histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p99' - - expr: sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-avg' - title: Flow Operation Latency type: timeseries description: Flow Operation Latency. @@ -2232,24 +2208,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-p99' - - title: Flow Buffer Size per Instance - type: timeseries - description: Flow Buffer Size per Instance. - queries: - - expr: greptime_flow_input_buf_size - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]' - - title: Flow Processing Error per Instance - type: timeseries - description: Flow Processing Error per Instance. - queries: - - expr: sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-[{{code}}]' - title: Flow Batching Attempt / Error Rate type: timeseries description: Flow batching mode task attempt, error, and slow-query rates by flow. diff --git a/grafana/dashboards/metrics/standalone/dashboard.json b/grafana/dashboards/metrics/standalone/dashboard.json index a90d05760e4..c1fda955d90 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.json +++ b/grafana/dashboards/metrics/standalone/dashboard.json @@ -3585,19 +3585,6 @@ }, "pluginVersion": "11.6.0", "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by (code) (rate(greptime_flow_errors[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "flow-{{code}}", - "range": true, - "refId": "A" - }, { "datasource": { "type": "prometheus", @@ -11583,7 +11570,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "hit-{{type}}", @@ -11596,7 +11583,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "miss-{{type}}", @@ -11609,7 +11596,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "evict-{{type}}-{{cause}}", @@ -16622,127 +16609,6 @@ "title": "Flow Ingest / Output Rate", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Ingest Latency.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "points", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 170 - }, - "id": 521, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p95", - "range": true, - "refId": "A" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p99", - "range": true, - "refId": "B" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-avg", - "range": true, - "refId": "C" - } - ], - "title": "Flow Ingest Latency", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -16805,7 +16671,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, + "x": 8, "y": 170 }, "id": 522, @@ -16851,196 +16717,6 @@ "title": "Flow Operation Latency", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Buffer Size per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 0, - "y": 178 - }, - "id": 523, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "greptime_flow_input_buf_size", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Buffer Size per Instance", - "type": "timeseries" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Processing Error per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 178 - }, - "id": 524, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-[{{code}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Processing Error per Instance", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -17105,7 +16781,7 @@ "h": 8, "w": 8, "x": 16, - "y": 178 + "y": 170 }, "id": 558, "options": { @@ -17229,7 +16905,7 @@ "h": 8, "w": 8, "x": 0, - "y": 186 + "y": 178 }, "id": 559, "options": { @@ -17353,7 +17029,7 @@ "h": 8, "w": 8, "x": 8, - "y": 186 + "y": 178 }, "id": 560, "options": { @@ -17464,7 +17140,7 @@ "h": 8, "w": 8, "x": 16, - "y": 186 + "y": 178 }, "id": 561, "options": { @@ -17575,7 +17251,7 @@ "h": 8, "w": 8, "x": 0, - "y": 194 + "y": 186 }, "id": 562, "options": { @@ -17673,7 +17349,7 @@ "h": 8, "w": 8, "x": 8, - "y": 194 + "y": 186 }, "id": 563, "options": { @@ -17771,7 +17447,7 @@ "h": 8, "w": 8, "x": 16, - "y": 194 + "y": 186 }, "id": 564, "options": { diff --git a/grafana/dashboards/metrics/standalone/dashboard.md b/grafana/dashboards/metrics/standalone/dashboard.md index 2ed84c59213..5dd6f1600e1 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.md +++ b/grafana/dashboards/metrics/standalone/dashboard.md @@ -36,7 +36,7 @@ | Scan and Compaction Memory Rejects | `sum(rate(greptime_mito_scan_requests_rejected_total[$__rate_interval]))`
`sum(rate(greptime_mito_scan_memory_exhausted_total[$__rate_interval]))`
`sum(rate(greptime_mito_compaction_memory_rejected_total[$__rate_interval]))` | `timeseries` | Datanode scan and compaction memory rejection/exhaustion counters. | `prometheus` | `rps` | `scan-rejected` | | OpenDAL Errors | `sum by (scheme, operation, error) (rate(opendal_operation_errors_total{error!="NotFound"}[$__rate_interval]))` | `timeseries` | Object-store errors by scheme, operation, and error, excluding NotFound noise. | `prometheus` | `eps` | `{{scheme}}-{{operation}}-{{error}}` | | Metasrv Failures | `sum(rate(greptime_meta_region_migration_fail[$__rate_interval]))`
`sum(rate(greptime_meta_reconciliation_procedure_error[$__rate_interval]))` | `timeseries` | Region migration and reconciliation failures in metasrv. | `prometheus` | `eps` | `migration-fail` | -| Flow and Trigger Failures | `sum by (code) (rate(greptime_flow_errors[$__rate_interval]))`
`sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `flow-{{code}}` | +| Flow and Trigger Failures | `sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `trigger-evaluate` | | Mito GC Failures | `sum(rate(greptime_mito_gc_errors_total[$__rate_interval]))`
`sum(rate(greptime_mito_gc_orphaned_index_files[$__rate_interval]))`
`sum(rate(greptime_mito_gc_skipped_unparsable_files[$__rate_interval]))` | `timeseries` | Mito garbage-collection errors and skipped/orphaned files on datanodes. | `prometheus` | `short` | `gc-errors` | # Capacity | Title | Query | Type | Description | Datasource | Unit | Legend Format | @@ -131,7 +131,7 @@ | Index Memory Usage | `greptime_index_apply_memory_usage`
`sum by (type) (greptime_index_create_memory_usage)` | `timeseries` | Memory used while applying and creating indexes. Growth here can explain memory pressure during indexed flush or compaction work. | `prometheus` | `bytes` | `apply` | | Index IO Bytes | `sum by (type, file_type) (rate(greptime_index_io_bytes_total[$__rate_interval]))` | `timeseries` | Index read and write byte rates by operation and file type for puffin and intermediate files. | `prometheus` | `Bps` | `{{type}}-{{file_type}}` | | Index IO Operations | `sum by (type, file_type) (rate(greptime_index_io_op_total[$__rate_interval]))` | `timeseries` | Index IO operation rates by operation and file type, including read, write, seek, and flush operations. | `prometheus` | `ops` | `{{type}}-{{file_type}}` | -| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | +| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | # Metasrv | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | @@ -278,10 +278,7 @@ ORDER BY data_size DESC;` | `piechart` | Distribution of leader regions and data | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | | Flow Ingest / Output Rate | `sum by(instance, pod, direction) (rate(greptime_flow_processed_rows[$__rate_interval]))` | `timeseries` | Flow Ingest / Output Rate. | `prometheus` | -- | `[{{pod}}]-[{{instance}}]-[{{direction}}]` | -| Flow Ingest Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))` | `timeseries` | Flow Ingest Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-p95` | | Flow Operation Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))`
`histogram_quantile(0.99, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))` | `timeseries` | Flow Operation Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{type}}]-p95` | -| Flow Buffer Size per Instance | `greptime_flow_input_buf_size` | `timeseries` | Flow Buffer Size per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]` | -| Flow Processing Error per Instance | `sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))` | `timeseries` | Flow Processing Error per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{code}}]` | | Flow Batching Attempt / Error Rate | `sum by(instance,pod,flow_id) (rate(greptime_flow_batching_start_query_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_error_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_slow_query_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode task attempt, error, and slow-query rates by flow. | `prometheus` | `ops` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-attempt` | | Flow Batching Query Latency | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`histogram_quantile(0.99, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode query latency by flow. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | | Flow Batching Query Window Count | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_count[$__rate_interval]))` | `timeseries` | Flow batching mode query window count by flow. | `prometheus` | `short` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | diff --git a/grafana/dashboards/metrics/standalone/dashboard.yaml b/grafana/dashboards/metrics/standalone/dashboard.yaml index 8da2a8fa6f1..beadb9c3215 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.yaml +++ b/grafana/dashboards/metrics/standalone/dashboard.yaml @@ -515,11 +515,6 @@ groups: description: Derived-data and alerting pipeline failures. unit: eps queries: - - expr: sum by (code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: flow-{{code}} - expr: sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval])) datasource: type: prometheus @@ -1572,17 +1567,17 @@ groups: description: Index-related cache hits, misses, and evictions from Mito caches. unit: ops queries: - - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: hit-{{type}} - - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: miss-{{type}} - - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} @@ -2199,25 +2194,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{pod}}]-[{{instance}}]-[{{direction}}]' - - title: Flow Ingest Latency - type: timeseries - description: Flow Ingest Latency. - queries: - - expr: histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p95' - - expr: histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p99' - - expr: sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-avg' - title: Flow Operation Latency type: timeseries description: Flow Operation Latency. @@ -2232,24 +2208,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-p99' - - title: Flow Buffer Size per Instance - type: timeseries - description: Flow Buffer Size per Instance. - queries: - - expr: greptime_flow_input_buf_size - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]' - - title: Flow Processing Error per Instance - type: timeseries - description: Flow Processing Error per Instance. - queries: - - expr: sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-[{{code}}]' - title: Flow Batching Attempt / Error Rate type: timeseries description: Flow batching mode task attempt, error, and slow-query rates by flow. diff --git a/rust-toolchain.toml b/rust-toolchain.toml index d16edecca80..642220914d8 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2026-03-21" +channel = "1.98.1" diff --git a/rustfmt.toml b/rustfmt.toml index 64d94def266..fccaefbc69a 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,2 +1,3 @@ -group_imports = "StdExternalCrate" -imports_granularity = "Module" +# disable these features until they become stable +#group_imports = "StdExternalCrate" +#imports_granularity = "Module" diff --git a/scripts/check-builder-rust-version.sh b/scripts/check-builder-rust-version.sh index fcfbd7d189d..3ca7acf9f84 100755 --- a/scripts/check-builder-rust-version.sh +++ b/scripts/check-builder-rust-version.sh @@ -16,26 +16,50 @@ function check_rust_toolchain_version() { DEV_BUILDER_UBUNTU_IMAGE="$DEV_BUILDER_UBUNTU_REGISTRY/$DEV_BUILDER_UBUNTU_NAMESPACE/$DEV_BUILDER_UBUNTU_NAME:$DEV_BUILDER_IMAGE_TAG" - CURRENT_VERSION=$(grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}' "$RUST_TOOLCHAIN_VERSION_FILE") - if [ -z "$CURRENT_VERSION" ]; then - echo "Error: No rust toolchain version found in $RUST_TOOLCHAIN_VERSION_FILE" + # The pinned channel, e.g. "1.96.1" or "nightly-2026-03-21". + RUST_TOOLCHAIN=$(grep -E '^channel' "$RUST_TOOLCHAIN_VERSION_FILE" | cut -d'"' -f2) + if [ -z "$RUST_TOOLCHAIN" ]; then + echo "Error: No rust toolchain channel found in $RUST_TOOLCHAIN_VERSION_FILE" exit 1 fi - RUST_TOOLCHAIN_VERSION_IN_BUILDER=$(docker run "$DEV_BUILDER_UBUNTU_IMAGE" rustc --version | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') - if [ -z "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" ]; then + RUSTC_VERSION_OUTPUT=$(docker run "$DEV_BUILDER_UBUNTU_IMAGE" rustc --version) + RUSTC_VERSION=$(echo "$RUSTC_VERSION_OUTPUT" | awk '{print $2}') + if [ -z "$RUSTC_VERSION" ]; then echo "Error: No rustc version found in $DEV_BUILDER_UBUNTU_IMAGE" exit 1 fi - # Compare the version and the difference should be less than 1 day. - current_rust_toolchain_seconds=$(date -d "$CURRENT_VERSION" +%s) - rust_toolchain_in_dev_builder_ubuntu_seconds=$(date -d "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" +%s) - date_diff=$(( (current_rust_toolchain_seconds - rust_toolchain_in_dev_builder_ubuntu_seconds) / 86400 )) + if [[ "$RUST_TOOLCHAIN" =~ ^[0-9]+\.[0-9]+ ]]; then + # Stable channel: the builder image must ship exactly the pinned version. + if [ "$RUSTC_VERSION" != "$RUST_TOOLCHAIN" ]; then + echo "Error: The rust toolchain '$RUSTC_VERSION' in builder '$DEV_BUILDER_UBUNTU_IMAGE' doesn't match the pinned '$RUST_TOOLCHAIN', please rebuild the dev-builder image" + exit 1 + fi + else + # Nightly channel: keep the legacy date-based check. The difference + # between the pinned nightly date and the rustc build date shipped in + # the builder should be less than 1 day. + CURRENT_VERSION=$(echo "$RUST_TOOLCHAIN" | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') + if [ -z "$CURRENT_VERSION" ]; then + echo "Error: Unsupported toolchain channel '$RUST_TOOLCHAIN' in $RUST_TOOLCHAIN_VERSION_FILE" + exit 1 + fi - if [ $date_diff -gt 1 ]; then - echo "Error: The rust toolchain '$RUST_TOOLCHAIN_VERSION_IN_BUILDER' in builder '$DEV_BUILDER_UBUNTU_IMAGE' maybe outdated, please update it to '$CURRENT_VERSION'" - exit 1 + RUST_TOOLCHAIN_VERSION_IN_BUILDER=$(echo "$RUSTC_VERSION_OUTPUT" | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') + if [ -z "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" ]; then + echo "Error: No rustc version found in $DEV_BUILDER_UBUNTU_IMAGE" + exit 1 + fi + + current_rust_toolchain_seconds=$(date -d "$CURRENT_VERSION" +%s) + rust_toolchain_in_dev_builder_ubuntu_seconds=$(date -d "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" +%s) + date_diff=$(( (current_rust_toolchain_seconds - rust_toolchain_in_dev_builder_ubuntu_seconds) / 86400 )) + + if [ $date_diff -gt 1 ]; then + echo "Error: The rust toolchain '$RUST_TOOLCHAIN_VERSION_IN_BUILDER' in builder '$DEV_BUILDER_UBUNTU_IMAGE' maybe outdated, please update it to '$CURRENT_VERSION'" + exit 1 + fi fi } diff --git a/scripts/fix-udeps.py b/scripts/fix-udeps.py deleted file mode 100755 index 7547b897361..00000000000 --- a/scripts/fix-udeps.py +++ /dev/null @@ -1,265 +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. - -import json -import os -import re -import sys - - -def load_udeps_report(report_path): - try: - with open(report_path, "r") as f: - return json.load(f) - except FileNotFoundError: - print(f"Error: Report file '{report_path}' not found.") - return None - except json.JSONDecodeError as e: - print(f"Error: Invalid JSON in report file: {e}") - return None - - -def extract_unused_dependencies(report): - """ - Extract and organize unused dependencies from the cargo-udeps JSON report. - - The cargo-udeps report has this structure: - { - "unused_deps": { - "package_name v0.1.0 (/path/to/package)": { - "normal": ["dep1", "dep2"], - "development": ["dev_dep1"], - "build": ["build_dep1"], - "manifest_path": "/path/to/Cargo.toml" - } - } - } - - Args: - report (dict): The parsed JSON report from cargo-udeps - - Returns: - dict: Organized unused dependencies by package name: - { - "package_name": { - "dependencies": [("dep1", "normal"), ("dev_dep1", "dev")], - "manifest_path": "/path/to/Cargo.toml" - } - } - """ - if not report or "unused_deps" not in report: - return {} - - unused_deps = {} - for package_full_name, deps_info in report["unused_deps"].items(): - package_name = package_full_name.split(" ")[0] - - all_unused = [] - if deps_info.get("normal"): - all_unused.extend([(dep, "normal") for dep in deps_info["normal"]]) - if deps_info.get("development"): - all_unused.extend([(dep, "dev") for dep in deps_info["development"]]) - if deps_info.get("build"): - all_unused.extend([(dep, "build") for dep in deps_info["build"]]) - - if all_unused: - unused_deps[package_name] = { - "dependencies": all_unused, - "manifest_path": deps_info.get("manifest_path", "unknown"), - } - - return unused_deps - - -def get_section_pattern(dep_type): - """ - Get regex patterns to identify different dependency sections in Cargo.toml. - - Args: - dep_type (str): Type of dependency ("normal", "dev", or "build") - - Returns: - list: List of regex patterns to match the appropriate section headers - - """ - patterns = { - "normal": [r"\[dependencies\]", r"\[dependencies\..*?\]"], - "dev": [r"\[dev-dependencies\]", r"\[dev-dependencies\..*?\]"], - "build": [r"\[build-dependencies\]", r"\[build-dependencies\..*?\]"], - } - return patterns.get(dep_type, []) - - -def remove_dependency_line(content, dep_name, section_start, section_end): - """ - Remove a dependency line from a specific section of a Cargo.toml file. - - Args: - content (str): The entire content of the Cargo.toml file - dep_name (str): Name of the dependency to remove (e.g., "serde", "tokio") - section_start (int): Starting position of the section in the content - section_end (int): Ending position of the section in the content - - Returns: - tuple: (new_content, removed) where: - - new_content (str): The modified content with dependency removed - - removed (bool): True if dependency was found and removed, False otherwise - - Example input content format: - content = ''' - [package] - name = "my-crate" - version = "0.1.0" - - [dependencies] - serde = "1.0" - tokio = { version = "1.0", features = ["full"] } - serde_json.workspace = true - - [dev-dependencies] - tempfile = "3.0" - ''' - - # If dep_name = "serde", section_start = start of [dependencies], - # section_end = start of [dev-dependencies], this function will: - # 1. Extract the section: "serde = "1.0"\ntokio = { version = "1.0", features = ["full"] }\nserde_json.workspace = true\n" - # 2. Find and remove the line: "serde = "1.0"" - # 3. Return the modified content with that line removed - """ - section_content = content[section_start:section_end] - - dep_patterns = [ - rf"^{re.escape(dep_name)}\s*=.*$", # e.g., "serde = "1.0"" - rf"^{re.escape(dep_name)}\.workspace\s*=.*$", # e.g., "serde_json.workspace = true" - ] - - for pattern in dep_patterns: - match = re.search(pattern, section_content, re.MULTILINE) - if match: - line_start = section_start + match.start() # Start of the matched line - line_end = section_start + match.end() # End of the matched line - - if line_end < len(content) and content[line_end] == "\n": - line_end += 1 - - return content[:line_start] + content[line_end:], True - - return content, False - - -def remove_dependency_from_toml(file_path, dep_name, dep_type): - """ - Remove a specific dependency from a Cargo.toml file. - - Args: - file_path (str): Path to the Cargo.toml file - dep_name (str): Name of the dependency to remove - dep_type (str): Type of dependency ("normal", "dev", or "build") - - Returns: - bool: True if dependency was successfully removed, False otherwise - """ - try: - with open(file_path, "r") as f: - content = f.read() - - section_patterns = get_section_pattern(dep_type) - if not section_patterns: - return False - - for pattern in section_patterns: - section_match = re.search(pattern, content, re.IGNORECASE) - if not section_match: - continue - - section_start = section_match.end() - next_section = re.search(r"\n\s*\[", content[section_start:]) - section_end = ( - section_start + next_section.start() if next_section else len(content) - ) - - new_content, removed = remove_dependency_line( - content, dep_name, section_start, section_end - ) - if removed: - with open(file_path, "w") as f: - f.write(new_content) - return True - - return False - - except Exception as e: - print(f"Error processing {file_path}: {e}") - return False - - -def process_unused_dependencies(unused_deps): - """ - Process and remove all unused dependencies from their respective Cargo.toml files. - - Args: - unused_deps (dict): Dictionary of unused dependencies organized by package: - { - "package_name": { - "dependencies": [("dep1", "normal"), ("dev_dep1", "dev")], - "manifest_path": "/path/to/Cargo.toml" - } - } - - """ - if not unused_deps: - print("No unused dependencies found.") - return - - total_removed = 0 - total_failed = 0 - - for package, info in unused_deps.items(): - deps = info["dependencies"] - manifest_path = info["manifest_path"] - - if not os.path.exists(manifest_path): - print(f"Manifest file not found: {manifest_path}") - total_failed += len(deps) - continue - - for dep, dep_type in deps: - if remove_dependency_from_toml(manifest_path, dep, dep_type): - print(f"Removed {dep} from {package}") - total_removed += 1 - else: - print(f"Failed to remove {dep} from {package}") - total_failed += 1 - - print(f"Removed {total_removed} dependencies") - if total_failed > 0: - print(f"Failed to remove {total_failed} dependencies") - - -def main(): - if len(sys.argv) > 1: - report_path = sys.argv[1] - else: - report_path = "udeps-report.json" - - report = load_udeps_report(report_path) - if report is None: - sys.exit(1) - - unused_deps = extract_unused_dependencies(report) - process_unused_dependencies(unused_deps) - - -if __name__ == "__main__": - main() diff --git a/src/api/src/helper.rs b/src/api/src/helper.rs index 40f9bdfb6d6..c23ae2587d5 100644 --- a/src/api/src/helper.rs +++ b/src/api/src/helper.rs @@ -87,6 +87,28 @@ impl ColumnDataTypeWrapper { } } +/// Returns the time unit if `datatype` is a timestamp type. +pub fn timestamp_unit(datatype: ColumnDataType) -> Option { + match datatype { + ColumnDataType::TimestampSecond => Some(TimeUnit::Second), + ColumnDataType::TimestampMillisecond => Some(TimeUnit::Millisecond), + ColumnDataType::TimestampMicrosecond => Some(TimeUnit::Microsecond), + ColumnDataType::TimestampNanosecond => Some(TimeUnit::Nanosecond), + _ => None, + } +} + +/// Returns the timestamp [ColumnDataType] for the given time unit. +/// This is the inverse of [timestamp_unit]. +pub fn timestamp_datatype(unit: TimeUnit) -> ColumnDataType { + match unit { + TimeUnit::Second => ColumnDataType::TimestampSecond, + TimeUnit::Millisecond => ColumnDataType::TimestampMillisecond, + TimeUnit::Microsecond => ColumnDataType::TimestampMicrosecond, + TimeUnit::Nanosecond => ColumnDataType::TimestampNanosecond, + } +} + impl From for ConcreteDataType { fn from(datatype_wrapper: ColumnDataTypeWrapper) -> Self { match datatype_wrapper.datatype { @@ -1215,6 +1237,21 @@ mod tests { use super::*; use crate::v1::Column; + #[test] + fn test_timestamp_unit_roundtrip() { + for unit in [ + TimeUnit::Second, + TimeUnit::Millisecond, + TimeUnit::Microsecond, + TimeUnit::Nanosecond, + ] { + assert_eq!(timestamp_unit(timestamp_datatype(unit)), Some(unit)); + } + // Non-timestamp types have no time unit. + assert_eq!(timestamp_unit(ColumnDataType::String), None); + assert_eq!(timestamp_unit(ColumnDataType::Datetime), None); + } + #[test] fn test_values_with_capacity() { let values = values_with_capacity(ColumnDataType::Int8, 2); diff --git a/src/api/src/v1/column_def.rs b/src/api/src/v1/column_def.rs index 9b39d0f3296..b02b06feafd 100644 --- a/src/api/src/v1/column_def.rs +++ b/src/api/src/v1/column_def.rs @@ -18,7 +18,7 @@ use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_K use datatypes::schema::{ COMMENT_KEY, ColumnDefaultConstraint, ColumnSchema, FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, - SkippingIndexOptions, SkippingIndexType, VECTOR_INDEX_KEY, + SkippingIndexOptions, SkippingIndexType, }; use greptime_proto::v1::{ Analyzer, FulltextBackend as PbFulltextBackend, SkippingIndexType as PbSkippingIndexType, @@ -35,14 +35,11 @@ const FULLTEXT_GRPC_KEY: &str = "fulltext"; const INVERTED_INDEX_GRPC_KEY: &str = "inverted_index"; /// Key used to store skip index options in gRPC column options. const SKIPPING_INDEX_GRPC_KEY: &str = "skipping_index"; -/// Key used to store vector index options in gRPC column options. -const VECTOR_INDEX_GRPC_KEY: &str = "vector_index"; -const COLUMN_OPTION_MAPPINGS: [(&str, &str); 6] = [ +const COLUMN_OPTION_MAPPINGS: [(&str, &str); 5] = [ (FULLTEXT_GRPC_KEY, FULLTEXT_KEY), (INVERTED_INDEX_GRPC_KEY, INVERTED_INDEX_KEY), (SKIPPING_INDEX_GRPC_KEY, SKIPPING_INDEX_KEY), - (VECTOR_INDEX_GRPC_KEY, VECTOR_INDEX_KEY), (EXTENSION_TYPE_NAME_KEY, EXTENSION_TYPE_NAME_KEY), (EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_METADATA_KEY), ]; @@ -80,9 +77,6 @@ pub fn try_as_column_schema(column_def: &ColumnDef) -> Result { if let Some(skipping_index) = options.options.get(SKIPPING_INDEX_GRPC_KEY) { metadata.insert(SKIPPING_INDEX_KEY.to_string(), skipping_index.to_owned()); } - if let Some(vector_index) = options.options.get(VECTOR_INDEX_GRPC_KEY) { - metadata.insert(VECTOR_INDEX_KEY.to_string(), vector_index.to_owned()); - } if let Some(extension_name) = options.options.get(EXTENSION_TYPE_NAME_KEY) { metadata.insert(EXTENSION_TYPE_NAME_KEY.to_string(), extension_name.clone()); } @@ -178,11 +172,6 @@ pub fn options_from_column_schema(column_schema: &ColumnSchema) -> Option Skipp mod tests { use datatypes::data_type::ConcreteDataType; - use datatypes::schema::{ - FulltextAnalyzer, FulltextBackend, VectorDistanceMetric, VectorIndexOptions, - }; - use serde_json::json; + use datatypes::schema::{FulltextAnalyzer, FulltextBackend}; use super::*; use crate::v1::ColumnDataType; @@ -297,10 +283,6 @@ mod tests { "{\"enable\":true}".to_string(), ), (INVERTED_INDEX_GRPC_KEY.to_string(), "true".to_string()), - ( - VECTOR_INDEX_GRPC_KEY.to_string(), - "{\"engine\":\"usearch\",\"metric\":\"l2sq\",\"connectivity\":16,\"expansion-add\":128,\"expansion-search\":64}".to_string(), - ), ]), }), }; @@ -323,8 +305,6 @@ mod tests { } ); assert!(schema.is_inverted_indexed()); - let vector_options = schema.vector_index_options().unwrap().unwrap(); - assert_eq!(vector_options.metric, VectorDistanceMetric::L2sq); } #[test] @@ -355,29 +335,6 @@ mod tests { ); } - #[test] - fn test_vector_index_options_roundtrip() { - let schema = ColumnSchema::new("test", ConcreteDataType::vector_datatype(4), true) - .with_vector_index_options(&VectorIndexOptions::default()) - .unwrap(); - let column_def = try_as_column_def(&schema, false).unwrap(); - let roundtrip = try_as_column_schema(&column_def).unwrap(); - let options = roundtrip.vector_index_options().unwrap().unwrap(); - assert_eq!(options.metric, VectorDistanceMetric::L2sq); - - let options = column_def.options.unwrap(); - let raw = options.options.get(VECTOR_INDEX_GRPC_KEY).unwrap(); - let json_value: serde_json::Value = serde_json::from_str(raw).unwrap(); - let expected = json!({ - "engine": "usearch", - "metric": "l2sq", - "connectivity": 16, - "expansion-add": 128, - "expansion-search": 64 - }); - assert_eq!(json_value, expected); - } - #[test] fn test_options_with_fulltext() { let fulltext = FulltextOptions::new_unchecked( diff --git a/src/auth/src/user_provider/watch_file_user_provider.rs b/src/auth/src/user_provider/watch_file_user_provider.rs index 3208948b11d..bba97f6df1f 100644 --- a/src/auth/src/user_provider/watch_file_user_provider.rs +++ b/src/auth/src/user_provider/watch_file_user_provider.rs @@ -199,4 +199,46 @@ pub mod test { test_authenticate(&provider, "root", "654321", false, Some(timeout)).await; test_authenticate(&provider, "admin", "654321", false, Some(timeout)).await; } + + /// Regression test for #9310: the users file is reached through a chain + /// of symlinks and a directory symlink in the chain gets swapped + /// atomically (like `ln -sfn`, as done by sops-nix). + #[cfg(unix)] + #[tokio::test] + async fn test_file_provider_symlink_chain() { + use std::os::unix::fs::symlink; + + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_provider_symlink_chain"); + let base = dir.path(); + let timeout = Duration::from_secs(10); + + // secrets.1/users <- secrets (dir symlink) <- conf/users (file symlink) + std::fs::create_dir(base.join("secrets.1")).unwrap(); + std::fs::write(base.join("secrets.1/users"), "root=123456\n").unwrap(); + symlink(base.join("secrets.1"), base.join("secrets")).unwrap(); + std::fs::create_dir(base.join("conf")).unwrap(); + symlink(base.join("secrets/users"), base.join("conf/users")).unwrap(); + + let file_path = base.join("conf/users"); + let provider = WatchFileUserProvider::new(file_path.to_str().unwrap()).unwrap(); + test_authenticate(&provider, "root", "123456", true, None).await; + + // Atomically swap the directory symlink to a new generation. + std::fs::create_dir(base.join("secrets.2")).unwrap(); + std::fs::write(base.join("secrets.2/users"), "root=654321\n").unwrap(); + symlink(base.join("secrets.2"), base.join("secrets.tmp")).unwrap(); + std::fs::rename(base.join("secrets.tmp"), base.join("secrets")).unwrap(); + std::fs::remove_dir_all(base.join("secrets.1")).unwrap(); + + test_authenticate(&provider, "root", "654321", true, Some(timeout)).await; + test_authenticate(&provider, "root", "123456", false, None).await; + + // In-place edits of the new target are picked up too, which requires + // the watch to follow the chain to the new directory. + std::fs::write(base.join("secrets.2/users"), "root=abcdef\n").unwrap(); + test_authenticate(&provider, "root", "abcdef", true, Some(timeout)).await; + test_authenticate(&provider, "root", "654321", false, None).await; + } } diff --git a/src/cache/src/lib.rs b/src/cache/src/lib.rs index 1a234aa802e..c7304ae6009 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -20,8 +20,8 @@ use std::time::Duration; use catalog::kvbackend::new_table_cache; use common_meta::cache::{ - CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistryBuilder, new_schema_cache, - new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache, + CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder, + new_schema_cache, new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache, new_table_route_cache, new_table_schema_cache, new_view_info_cache, }; use common_meta::kv_backend::KvBackendRef; @@ -53,11 +53,11 @@ pub const TABLE_FLOWNODE_SET_CACHE_NAME: &str = "table_flownode_set_cache"; pub const TABLE_ROUTE_CACHE_NAME: &str = "table_route_cache"; pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache"; -/// Builds cache registry for datanode, including: -/// - Schema cache. -/// - Table id to schema name cache. -pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry { - // Builds table id schema name cache that never expires. +/// Builds the datanode cache registry with base metadata caches before derived caches. +/// +/// Invalidating the base layer first prevents a concurrent query from refilling a derived cache +/// from stale metadata after that derived cache has been invalidated. +pub fn build_datanode_layered_cache_registry(kv_backend: KvBackendRef) -> LayeredCacheRegistry { let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build(); let table_id_schema_cache = Arc::new(new_table_schema_cache( TABLE_SCHEMA_NAME_CACHE_NAME.to_string(), @@ -65,7 +65,6 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry kv_backend.clone(), )); - // Builds schema cache let cache = default_cache(); let schema_cache = Arc::new(new_schema_cache( SCHEMA_CACHE_NAME.to_string(), @@ -73,9 +72,57 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry kv_backend.clone(), )); - CacheRegistryBuilder::default() + let cache = default_cache(); + let table_info_cache = Arc::new(new_table_info_cache( + TABLE_INFO_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_name_cache = Arc::new(new_table_name_cache( + TABLE_NAME_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_route_cache = Arc::new(new_table_route_cache( + TABLE_ROUTE_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_cache = Arc::new(new_table_cache( + TABLE_CACHE_NAME.to_string(), + cache, + table_info_cache.clone(), + table_name_cache.clone(), + )); + + let cache = default_cache(); + let partition_info_cache = Arc::new(new_partition_info_cache( + PARTITION_INFO_CACHE_NAME.to_string(), + cache, + table_route_cache.clone(), + )); + + let base_registry = CacheRegistryBuilder::default() .add_cache(table_id_schema_cache) .add_cache(schema_cache) + .add_cache(table_info_cache) + .add_cache(table_name_cache) + .add_cache(table_route_cache) + .build(); + let derived_registry = CacheRegistryBuilder::default() + .add_cache(table_cache) + .add_cache(partition_info_cache) + .build(); + + LayeredCacheRegistryBuilder::default() + .add_cache_registry(base_registry) + .add_cache_registry(derived_registry) .build() } diff --git a/src/catalog/Cargo.toml b/src/catalog/Cargo.toml index f46887e148a..fe42ddc2d13 100644 --- a/src/catalog/Cargo.toml +++ b/src/catalog/Cargo.toml @@ -61,9 +61,7 @@ tokio-stream.workspace = true [dev-dependencies] cache.workspace = true -catalog = { workspace = true, features = ["testing"] } chrono.workspace = true common-meta = { workspace = true, features = ["testing"] } common-query = { workspace = true, features = ["testing"] } -object-store.workspace = true tokio.workspace = true diff --git a/src/catalog/src/kvbackend/client.rs b/src/catalog/src/kvbackend/client.rs index 69b31d62322..402af1d9865 100644 --- a/src/catalog/src/kvbackend/client.rs +++ b/src/catalog/src/kvbackend/client.rs @@ -192,6 +192,10 @@ impl KvBackend for CachedKvBackend { } } + if miss_keys.is_empty() { + return Ok(BatchGetResponse { kvs }); + } + let batch_get_req = BatchGetRequest::new().with_keys(miss_keys.clone()); let pre_version = self.version(); @@ -503,6 +507,7 @@ mod tests { pub struct SimpleKvBackend { inner_map: DashMap, Vec>, get_execute_times: Arc, + batch_get_execute_times: Arc, } impl TxnService for SimpleKvBackend { @@ -520,6 +525,7 @@ mod tests { } async fn batch_get(&self, req: BatchGetRequest) -> Result { + self.batch_get_execute_times.fetch_add(1, Ordering::SeqCst); let mut kvs = Vec::with_capacity(req.keys.len()); for key in req.keys.iter() { if let Some(kv) = self.get(key).await? { @@ -571,6 +577,7 @@ mod tests { async fn test_cached_kv_backend() { let simple_kv = Arc::new(SimpleKvBackend::default()); let get_execute_times = simple_kv.get_execute_times.clone(); + let batch_get_execute_times = simple_kv.batch_get_execute_times.clone(); let cached_kv = CachedKvBackend::wrap(simple_kv); add_some_vals(&cached_kv).await; @@ -582,9 +589,11 @@ mod tests { assert_eq!(get_execute_times.load(Ordering::SeqCst), 0); for _ in 0..10 { - let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + assert_eq!(batch_get_resp.kvs.len(), 2); assert_eq!(get_execute_times.load(Ordering::SeqCst), 2); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 1); } let batch_get_req = BatchGetRequest { @@ -594,14 +603,29 @@ mod tests { let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); assert_eq!(get_execute_times.load(Ordering::SeqCst), 3); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2); for _ in 0..10 { - let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + assert_eq!(batch_get_resp.kvs.len(), 3); assert_eq!(get_execute_times.load(Ordering::SeqCst), 3); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2); } } + #[tokio::test] + async fn test_cached_kv_backend_empty_batch_get() { + let simple_kv = Arc::new(SimpleKvBackend::default()); + let cached_kv = CachedKvBackend::wrap(simple_kv.clone()); + + let response = cached_kv.batch_get(BatchGetRequest::new()).await.unwrap(); + + assert!(response.kvs.is_empty()); + assert_eq!(simple_kv.batch_get_execute_times.load(Ordering::SeqCst), 0); + assert_eq!(simple_kv.get_execute_times.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn test_cached_kv_backend_rejects_writes_with_read_only_inner() { let inner = Arc::new(MemoryKvBackend::::new()); diff --git a/src/catalog/src/system_schema/information_schema/partitions.rs b/src/catalog/src/system_schema/information_schema/partitions.rs index d2cce93150a..1bfe3f634fd 100644 --- a/src/catalog/src/system_schema/information_schema/partitions.rs +++ b/src/catalog/src/system_schema/information_schema/partitions.rs @@ -254,6 +254,51 @@ impl InformationSchemaPartitionsBuilder { let predicates = Predicates::from_scan_request(&request); + // Fast path: when the predicates pin both the schema and the table name + // with equality, resolve the table directly instead of enumerating all + // tables in all schemas and fetching partition metadata for each of them. + if let (Some(schema_name), Some(table_name)) = ( + predicates + .find_eq_value(TABLE_SCHEMA) + .and_then(Value::as_string), + predicates + .find_eq_value(TABLE_NAME) + .and_then(Value::as_string), + ) { + if let Some(table) = catalog_manager + .table(&catalog_name, &schema_name, &table_name, None) + .await? + { + let table_info = table.table_info(); + + // The enumeration path below skips temporary tables, keep the + // behavior identical. + if table_info.table_type != TableType::Temporary { + let mut table_partitions = partition_manager + .batch_find_table_partitions(&[table_info.table_id()]) + .await + .context(FindPartitionsSnafu)?; + + let partitions = table_partitions + .remove(&table_info.table_id()) + .unwrap_or_default(); + + self.add_partitions( + &predicates, + &table_info, + &catalog_name, + &schema_name, + &table_info.name, + &partitions, + ); + } + } + + // The pinned table fully determines the result: the table above was + // found (or not) unambiguously, so no other table can contribute rows. + return self.finish(); + } + for schema_name in catalog_manager.schema_names(&catalog_name, None).await? { let table_info_stream = catalog_manager .tables(&catalog_name, &schema_name, None) @@ -421,3 +466,253 @@ impl DfPartitionStream for InformationSchemaPartitions { )) } } + +#[cfg(test)] +mod tests { + use api::v1::meta::Peer; + use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, INFORMATION_SCHEMA_NAME}; + use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; + use common_meta::key::schema_name::SchemaNameKey; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::wal_provider::RegionWalOptions; + use datafusion::logical_expr::{BinaryExpr, Expr, Operator, col, lit}; + use futures::TryStreamExt; + use store_api::storage::RegionId; + + use super::*; + use crate::information_schema::NoopInformationExtension; + use crate::kvbackend::KvBackendCatalogManagerBuilder; + + const SCHEMA_1: &str = "partitions_schema_1"; + const SCHEMA_2: &str = "partitions_schema_2"; + const PARTITIONED_TABLE: &str = "partitioned_table"; + const MISSING_TABLE: &str = "missing_table"; + const TABLE_ID_1: TableId = 1; + const TABLE_ID_2: TableId = 2; + const REGION_NUMBER_1: u32 = 1; + const REGION_NUMBER_2: u32 = 2; + const REGION_NUMBER_7: u32 = 7; + + fn kv_catalog_manager() -> Arc { + let backend = Arc::new(MemoryKvBackend::default()); + let layered_cache_builder = LayeredCacheRegistryBuilder::default() + .add_cache_registry(CacheRegistryBuilder::default().build()); + let fundamental_cache_registry = build_fundamental_cache_registry(backend.clone()); + let layered_cache_registry = Arc::new( + with_default_composite_cache_registry( + layered_cache_builder.add_cache_registry(fundamental_cache_registry), + ) + .unwrap() + .build(), + ); + + KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + backend, + layered_cache_registry, + ) + .build() + } + + fn region_route(table_id: TableId, region_number: u32) -> RegionRoute { + RegionRoute { + region: Region { + id: RegionId::new(table_id, region_number), + ..Default::default() + }, + leader_peer: Some(Peer { + id: 1, + addr: "127.0.0.1:3001".to_string(), + }), + ..Default::default() + } + } + + fn table_info(schema: &str, table: &str, table_id: TableId) -> TableInfo { + // The catalog defaults to the default catalog name. + let mut table_info = + common_meta::key::test_utils::new_test_table_info_with_name(table_id, table); + table_info.schema_name = schema.to_string(); + table_info + } + + /// Creates a catalog with two partitioned tables named `PARTITIONED_TABLE` + /// in two different schemas: + /// + /// - `SCHEMA_1.PARTITIONED_TABLE` has two regions. + /// - `SCHEMA_2.PARTITIONED_TABLE` has one region. + async fn setup_catalog() -> Arc { + let manager = kv_catalog_manager(); + let metadata = manager.table_metadata_manager_ref().clone(); + + for schema in [SCHEMA_1, SCHEMA_2] { + metadata + .schema_manager() + .create(SchemaNameKey::new(DEFAULT_CATALOG_NAME, schema), None, true) + .await + .unwrap(); + } + + for (table_id, schema, region_numbers) in [ + (TABLE_ID_1, SCHEMA_1, vec![REGION_NUMBER_1, REGION_NUMBER_2]), + (TABLE_ID_2, SCHEMA_2, vec![REGION_NUMBER_7]), + ] { + let routes = region_numbers + .into_iter() + .map(|region_number| region_route(table_id, region_number)) + .collect(); + metadata + .create_table_metadata( + table_info(schema, PARTITIONED_TABLE, table_id), + TableRouteValue::physical(routes), + RegionWalOptions::default(), + ) + .await + .unwrap(); + } + + manager + } + + /// Scans `information_schema.partitions` and returns the total row count + /// and the pretty printed batches. + async fn scan_partitions( + manager: &Arc, + filters: Vec, + ) -> (usize, String) { + let table = manager + .table( + DEFAULT_CATALOG_NAME, + INFORMATION_SCHEMA_NAME, + PARTITIONS, + None, + ) + .await + .unwrap() + .unwrap(); + let stream = table + .scan_to_stream(ScanRequest { + filters, + ..Default::default() + }) + .await + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + let rows = batches.iter().map(|batch| batch.num_rows()).sum(); + let output = batches + .iter() + .map(|batch| batch.pretty_print()) + .collect::>() + .join("\n"); + (rows, output) + } + + fn eq_filter(column: &str, value: &str) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(column)), + Operator::Eq, + Box::new(lit(value)), + )) + } + + fn not_eq_filter(column: &str, value: &str) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(column)), + Operator::NotEq, + Box::new(lit(value)), + )) + } + + #[tokio::test] + async fn test_scan_pinned_table_skips_catalog_enumeration() { + let manager: Arc = setup_catalog().await; + + // No predicates: every partitioned table is listed. + let (rows, output) = scan_partitions(&manager, vec![]).await; + assert_eq!(3, rows, "{output}"); + + // Pins both the schema and the table name, so the scan resolves the + // table directly instead of enumerating all catalog entries. + let (rows, output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + eq_filter(TABLE_NAME, PARTITIONED_TABLE), + ], + ) + .await; + assert_eq!(2, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(!output.contains(SCHEMA_2), "{output}"); + assert!( + output.contains( + &RegionId::new(TABLE_ID_1, REGION_NUMBER_1) + .as_u64() + .to_string() + ), + "{output}" + ); + assert!( + output.contains( + &RegionId::new(TABLE_ID_1, REGION_NUMBER_2) + .as_u64() + .to_string() + ), + "{output}" + ); + + // The fallback enumeration path must return the identical result. The + // inequality on `table_name` doesn't pin a single name, so no direct + // lookup happens here. + let (fallback_rows, fallback_output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + not_eq_filter(TABLE_NAME, MISSING_TABLE), + ], + ) + .await; + assert_eq!(rows, fallback_rows, "{fallback_output}"); + assert_eq!(output, fallback_output); + + // Only the table name is pinned: the schema can't be resolved and both + // tables of that name are kept. + let (rows, output) = + scan_partitions(&manager, vec![eq_filter(TABLE_NAME, PARTITIONED_TABLE)]).await; + assert_eq!(3, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(output.contains(SCHEMA_2), "{output}"); + + // An unknown table in a pinned schema yields no rows. + let (rows, output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + eq_filter(TABLE_NAME, MISSING_TABLE), + ], + ) + .await; + assert_eq!(0, rows, "{output}"); + } + + #[tokio::test] + async fn test_scan_equality_under_disjunction_falls_back_to_enumeration() { + let manager: Arc = setup_catalog().await; + + // `table_schema = SCHEMA_1 OR table_name = PARTITIONED_TABLE` matches + // the tables in both schemas: equality nested under a disjunction must + // not pin a single table. + let filter = Expr::BinaryExpr(BinaryExpr::new( + Box::new(eq_filter(TABLE_SCHEMA, SCHEMA_1)), + Operator::Or, + Box::new(eq_filter(TABLE_NAME, PARTITIONED_TABLE)), + )); + let (rows, output) = scan_partitions(&manager, vec![filter]).await; + assert_eq!(3, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(output.contains(SCHEMA_2), "{output}"); + } +} diff --git a/src/catalog/src/system_schema/predicate.rs b/src/catalog/src/system_schema/predicate.rs index 7957ba454d4..cb71e51ed0e 100644 --- a/src/catalog/src/system_schema/predicate.rs +++ b/src/catalog/src/system_schema/predicate.rs @@ -214,6 +214,22 @@ impl Predicate { _ => None, } } + + /// Collects the values that constrain `column` to be equal to them at this + /// predicate's top-level conjunctions. + /// + /// Constraints under `Or`/`Not` are ignored because they don't guarantee + /// that `column` equals a single value. + fn collect_eq_values<'a>(&'a self, column: &str, values: &mut Vec<&'a Value>) { + match self { + Predicate::Eq(c, v) if c == column => values.push(v), + Predicate::And(left, right) => { + left.collect_eq_values(column, values); + right.collect_eq_values(column, values); + } + _ => {} + } + } } /// Perform SQL left LIKE right, return `None` if fail to evaluate. @@ -269,6 +285,25 @@ impl Predicates { } } + /// Returns the single equality value that constrains `column` in the + /// top-level conjunctions of the predicates. + /// + /// Returns `None` if there is no equality constraint on `column`, or if + /// the predicates constrain it with multiple distinct values. + pub fn find_eq_value(&self, column: &str) -> Option<&Value> { + let mut values = Vec::new(); + for predicate in &self.predicates { + predicate.collect_eq_values(column, &mut values); + } + + let (first, rest) = values.split_first()?; + if rest.iter().all(|v| v == first) { + Some(*first) + } else { + None + } + } + /// Evaluate the predicates with the row. /// returns true when all the predicates are satisfied or can't be evaluated. pub fn eval(&self, row: &[(&str, &Value)]) -> bool { @@ -556,6 +591,71 @@ mod tests { ); } + #[test] + fn test_find_eq_value() { + fn eq(name: &str, value: &str) -> Predicate { + Predicate::Eq(name.to_string(), Value::from(value)) + } + + fn and(left: Predicate, right: Predicate) -> Predicate { + Predicate::And(Box::new(left), Box::new(right)) + } + + fn or(left: Predicate, right: Predicate) -> Predicate { + Predicate::Or(Box::new(left), Box::new(right)) + } + + fn predicates(predicates: Vec) -> Predicates { + Predicates { predicates } + } + + // No predicates, no equality value. + assert!(predicates(vec![]).find_eq_value("a").is_none()); + + // No equality on the column. + let p = predicates(vec![Predicate::NotEq("a".to_string(), Value::from("x"))]); + assert!(p.find_eq_value("a").is_none()); + assert!(p.find_eq_value("b").is_none()); + + // Simple equality. + let p = predicates(vec![eq("a", "x")]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + assert!(p.find_eq_value("b").is_none()); + + // Duplicated equality on the same value. + let p = predicates(vec![and(eq("a", "x"), eq("a", "x"))]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + + // Conflicting equalities reject the lookup. + let p = predicates(vec![and(eq("a", "x"), eq("a", "y"))]); + assert!(p.find_eq_value("a").is_none()); + + // Conflicting equalities in separate predicates reject the lookup too. + let p = predicates(vec![eq("a", "x"), eq("a", "y")]); + assert!(p.find_eq_value("a").is_none()); + + // Equality nested in deeper conjunctions. + let p = predicates(vec![and(and(eq("b", "y"), eq("a", "x")), eq("a", "x"))]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + + // Equality under a disjunction or a negation doesn't constrain the column. + let p = predicates(vec![or(eq("a", "x"), eq("a", "x"))]); + assert!(p.find_eq_value("a").is_none()); + let p = predicates(vec![Predicate::Not(Box::new(eq("a", "x")))]); + assert!(p.find_eq_value("a").is_none()); + + // Equality value in a top-level conjunction composed of separate predicates. + let p = predicates(vec![ + eq("a", "x"), + Predicate::Like("b".to_string(), "%y".to_string(), false), + ]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + } + #[test] fn test_predicates_eval_row() { let wrong_row = [ diff --git a/src/cli/Cargo.toml b/src/cli/Cargo.toml index 587e5b8cffa..6abe504a3ab 100644 --- a/src/cli/Cargo.toml +++ b/src/cli/Cargo.toml @@ -18,10 +18,8 @@ workspace = true [dependencies] async-stream.workspace = true async-trait.workspace = true -auth.workspace = true base64.workspace = true cache.workspace = true -catalog.workspace = true chrono.workspace = true clap.workspace = true client = { workspace = true, features = ["testing"] } @@ -34,14 +32,11 @@ common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true common-procedure.workspace = true -common-query.workspace = true -common-recordbatch.workspace = true common-runtime.workspace = true common-telemetry = { workspace = true, features = [ "deadlock_detection", ] } common-time.workspace = true -common-version.workspace = true common-wal.workspace = true datatypes.workspace = true etcd-client.workspace = true @@ -61,21 +56,19 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true -session.workspace = true snafu.workspace = true sql.workspace = true standalone.workspace = true store-api.workspace = true table.workspace = true tokio.workspace = true -tracing-appender.workspace = true url.workspace = true uuid.workspace = true [dev-dependencies] common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true -common-version.workspace = true object-store = { workspace = true, features = ["testing"] } serde.workspace = true tempfile.workspace = true +tokio-util.workspace = true diff --git a/src/cli/src/common/store.rs b/src/cli/src/common/store.rs index e974037ea4d..c3b8e477ed1 100644 --- a/src/cli/src/common/store.rs +++ b/src/cli/src/common/store.rs @@ -86,7 +86,13 @@ pub struct StoreConfig { /// Automatically create PostgreSQL schema if it doesn't exist (default: true). #[cfg(feature = "pg_kvbackend")] - #[clap(long, default_value_t = true)] + #[clap( + long, + default_value_t = true, + default_missing_value = "true", + num_args = 0..=1, + action = clap::ArgAction::Set + )] pub auto_create_schema: bool, /// TLS mode for backend store connections (etcd, PostgreSQL, MySQL) @@ -223,3 +229,27 @@ impl StoreConfig { } } } + +#[cfg(all(test, feature = "pg_kvbackend"))] +mod tests { + use clap::Parser; + + use super::StoreConfig; + + #[test] + fn test_auto_create_schema_accepts_explicit_value() { + let config = + StoreConfig::try_parse_from(["store", "--auto-create-schema", "false"]).unwrap(); + assert!(!config.auto_create_schema); + + let config = + StoreConfig::try_parse_from(["store", "--auto-create-schema", "true"]).unwrap(); + assert!(config.auto_create_schema); + } + + #[test] + fn test_auto_create_schema_defaults_to_true() { + let config = StoreConfig::try_parse_from(["store"]).unwrap(); + assert!(config.auto_create_schema); + } +} diff --git a/src/cli/src/data.rs b/src/cli/src/data.rs index c880bd81c73..3e1924f12ee 100644 --- a/src/cli/src/data.rs +++ b/src/cli/src/data.rs @@ -18,6 +18,7 @@ mod import; pub mod import_v2; pub(crate) mod path; pub(crate) mod progress; +mod schema_export; pub mod snapshot_storage; pub(crate) mod sql; mod storage_export; diff --git a/src/cli/src/data/export.rs b/src/cli/src/data/export.rs index 148c27316e7..055c5d4ee02 100644 --- a/src/cli/src/data/export.rs +++ b/src/cli/src/data/export.rs @@ -27,6 +27,7 @@ use tokio::sync::Semaphore; use tokio::time::Instant; use crate::common::{ObjectStoreConfig, new_fs_object_store}; +use crate::data::schema_export::append_schema_ddl; use crate::data::storage_export::{ AzblobBackend, FsBackend, GcsBackend, OssBackend, S3Backend, StorageExport, StorageType, }; @@ -335,32 +336,6 @@ impl Export { )) } - async fn show_create( - &self, - show_type: &str, - catalog: &str, - schema: &str, - table: Option<&str>, - ) -> Result { - let sql = match table { - Some(table) => format!( - r#"SHOW CREATE {} "{}"."{}"."{}""#, - show_type, catalog, schema, table - ), - None => format!(r#"SHOW CREATE {} "{}"."{}""#, show_type, catalog, schema), - }; - let records = self - .database_client - .sql_in_public(&sql) - .await? - .context(EmptyResultSnafu)?; - let Value::String(create) = &records[0][1] else { - unreachable!() - }; - - Ok(format!("{};\n", create)) - } - async fn export_create_database(&self) -> Result<()> { let timer = Instant::now(); let db_names = self.get_db_names().await?; @@ -368,9 +343,15 @@ impl Export { let operator = self.build_prefer_fs_operator().await?; for schema in db_names { - let create_database = self - .show_create("DATABASE", &self.catalog, &schema, None) - .await?; + let mut create_database = String::new(); + append_schema_ddl( + &self.database_client, + &self.catalog, + &schema, + std::iter::once(("DATABASE", None)), + &mut create_database, + ) + .await?; let file_path = self.get_file_path(&schema, "create_database.sql"); self.write_to_storage(&operator, &file_path, create_database.into_bytes()) @@ -415,23 +396,28 @@ impl Export { } let file_path = export_self.get_file_path(&schema, "create_tables.sql"); - let mut content = Vec::new(); - - // Add table creation SQL - for (c, s, t) in metric_physical_tables.iter().chain(&remaining_tables) { - let create_table = export_self.show_create("TABLE", c, s, Some(t)).await?; - content.extend_from_slice(create_table.as_bytes()); - } - - // Add view creation SQL - for (c, s, v) in &views { - let create_view = export_self.show_create("VIEW", c, s, Some(v)).await?; - content.extend_from_slice(create_view.as_bytes()); - } + let mut content = String::new(); + let objects = metric_physical_tables + .iter() + .chain(&remaining_tables) + .map(|(_, _, table)| ("TABLE", Some(table.as_str()))) + .chain( + views + .iter() + .map(|(_, _, view)| ("VIEW", Some(view.as_str()))), + ); + append_schema_ddl( + &export_self.database_client, + &export_self.catalog, + &schema, + objects, + &mut content, + ) + .await?; // Write to storage export_self - .write_to_storage(&operator, &file_path, content) + .write_to_storage(&operator, &file_path, content.into_bytes()) .await?; info!( @@ -445,9 +431,11 @@ impl Export { }); } - let success = self.execute_tasks(tasks).await; + for result in futures::future::join_all(tasks).await { + result?; + } let elapsed = timer.elapsed(); - info!("Success {success}/{db_count} jobs, cost: {elapsed:?}"); + info!("Success {db_count}/{db_count} jobs, cost: {elapsed:?}"); Ok(()) } diff --git a/src/cli/src/data/export_v2.rs b/src/cli/src/data/export_v2.rs index f9a3c2be095..9887e525030 100644 --- a/src/cli/src/data/export_v2.rs +++ b/src/cli/src/data/export_v2.rs @@ -41,11 +41,13 @@ //! `--experimental-metric-export` enables shared Metric physical scans for Parquet. //! Enable `experimental_metric_export = true` on the frontend or standalone server //! and use an endpoint whose frontends all support and enable this option. +//! Add `--metric-data-layout packed` to write version-2 packed snapshots. +//! This requires the server packed-export capability; schema-only exports remain version 1. //! The command checks the server capability before modifying the snapshot. //! Each snapshot path belongs to one export task. Resume requires stable source //! schemas/data and confirmation that the previous export and storage writes ended; //! an HTTP timeout does not establish that. Only unfinished chunks are cleaned and -//! rerun. Completed chunks and the V2 snapshot/import format remain unchanged. +//! rerun. Resume retains the snapshot version/layout and completed chunks. mod chunker; mod command; diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index e0f67eb796c..48779f3c435 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -25,15 +25,15 @@ use common_error::ext::BoxedError; use common_telemetry::info; use serde_json::Value; use servers::http::{ColumnSchema, GreptimeQueryOutput, OutputSchema}; -use snafu::{OptionExt, ResultExt}; +use snafu::ResultExt; use crate::Tool; use crate::common::ObjectStoreConfig; use crate::data::export_v2::coordinator::{ExportDataOptions, export_data}; use crate::data::export_v2::error::{ - ChunkTimeWindowRequiresBoundsSnafu, DatabaseSnafu, EmptyResultSnafu, IoSnafu, - ManifestVersionMismatchSnafu, Result, ResumeConfigMismatchSnafu, SchemaOnlyArgsNotAllowedSnafu, - SchemaOnlyModeMismatchSnafu, SnapshotVerifyFailedSnafu, UnexpectedValueTypeSnafu, + ChunkTimeWindowRequiresBoundsSnafu, DatabaseSnafu, IoSnafu, ManifestVersionMismatchSnafu, + Result, ResumeConfigMismatchSnafu, SchemaOnlyArgsNotAllowedSnafu, SchemaOnlyModeMismatchSnafu, + SnapshotVerifyFailedSnafu, UnexpectedValueTypeSnafu, }; use crate::data::export_v2::extractor::SchemaExtractor; use crate::data::export_v2::manifest::{ @@ -42,10 +42,11 @@ use crate::data::export_v2::manifest::{ use crate::data::export_v2::schema::{DDL_DIR, SCHEMA_DIR, SCHEMAS_FILE}; use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; use crate::data::progress::{ProgressMode, build_progress_reporter}; +use crate::data::schema_export::append_schema_ddl; use crate::data::snapshot_storage::{ OpenDalStorage, SnapshotStorage, validate_snapshot_uri, validate_uri, }; -use crate::data::sql::{escape_sql_identifier, escape_sql_literal}; +use crate::data::sql::escape_sql_literal; use crate::database::{DatabaseClient, parse_proxy_opts}; /// Export V2 commands. @@ -296,6 +297,10 @@ pub struct ExportCreateCommand { #[clap(long)] experimental_metric_export: bool, + /// Store independent Metric Parquet streams in shared packed objects. + #[clap(long, value_parser = ["packed"], requires = "experimental_metric_export")] + metric_data_layout: Option, + /// Delete existing snapshot and recreate. #[clap(long)] force: bool, @@ -345,6 +350,16 @@ impl ExportCreateCommand { let time_range = TimeRange::parse(self.start_time.as_deref(), self.end_time.as_deref()) .map_err(BoxedError::new)?; + if self.metric_data_layout.is_some() + && time_range.is_bounded() + && time_range.start == time_range.end + { + return crate::error::InvalidArgumentsSnafu { + msg: "Packed export requires --start-time to be earlier than --end-time", + } + .fail() + .map_err(BoxedError::new); + } if self.chunk_time_window.is_some() && !time_range.is_bounded() { return ChunkTimeWindowRequiresBoundsSnafu .fail() @@ -425,6 +440,13 @@ impl ExportCreateCommand { .map_err(BoxedError::new); } } + if self.metric_data_layout.is_some() && !self.schema_only { + database_client + .require_packed_export() + .await + .context(DatabaseSnafu) + .map_err(BoxedError::new)?; + } let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?; Ok(Box::new(ExportCreate { @@ -434,6 +456,7 @@ impl ExportCreateCommand { schema_only: self.schema_only, format: self.format, experimental_metric_export: self.experimental_metric_export, + packed: self.metric_data_layout.is_some(), force: self.force, time_range, chunk_time_window: self.chunk_time_window, @@ -462,6 +485,7 @@ struct ExportConfig { schema_only: bool, format: DataFormat, experimental_metric_export: bool, + packed: bool, force: bool, time_range: TimeRange, chunk_time_window: Option, @@ -504,7 +528,7 @@ impl ExportCreate { let mut manifest = self.storage.read_manifest().await?; // Check version compatibility - if manifest.version != MANIFEST_VERSION || manifest.data_layout.is_some() { + if manifest.validate_layout().is_err() { return ManifestVersionMismatchSnafu { expected: MANIFEST_VERSION, found: manifest.version, @@ -513,6 +537,15 @@ impl ExportCreate { } validate_resume_config(&manifest, &self.config)?; + if manifest.is_packed() { + self.database_client + .require_packed_export() + .await + .context(DatabaseSnafu)?; + if !self.config.experimental_metric_export { + return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail(); + } + } info!( "Resuming existing snapshot: {} (completed: {}/{} chunks)", @@ -571,10 +604,14 @@ impl ExportCreate { self.config.chunk_time_window, )?; + if self.config.packed && !self.config.schema_only { + manifest.version = 2; + manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); + } if self.config.experimental_metric_export { for chunk in &manifest.chunks { self.storage - .prepare_export_chunk(&schema_names, chunk.id, false) + .prepare_export_chunk(&schema_names, chunk.id, false, manifest.is_packed()) .await?; } } @@ -584,8 +621,10 @@ impl ExportCreate { info!("Exported {} schemas", schema_snapshot.schemas.len()); // 5. Export DDL files for import recovery. - let ddl_by_schema = self.build_ddl_by_schema(&schema_names).await?; - for (schema, ddl) in ddl_by_schema { + let mut sorted_schemas = schema_names; + sorted_schemas.sort(); + for schema in sorted_schemas { + let ddl = self.build_schema_ddl(&schema).await?; let ddl_path = ddl_path_for_schema(&schema); self.storage.write_text(&ddl_path, &ddl).await?; info!("Exported DDL for schema {} to {}", schema, ddl_path); @@ -623,45 +662,31 @@ impl ExportCreate { Ok(()) } - async fn build_ddl_by_schema(&self, schema_names: &[String]) -> Result> { - let mut schemas = schema_names.to_vec(); - schemas.sort(); - - let mut ddl_by_schema = Vec::with_capacity(schemas.len()); - for schema in schemas { - let create_database = self.show_create("DATABASE", &schema, None).await?; - - let (mut physical_tables, mut tables, mut views) = - self.get_schema_objects(&schema).await?; - physical_tables.sort(); - let mut physical_ddls = Vec::with_capacity(physical_tables.len()); - for table in physical_tables { - physical_ddls.push(self.show_create("TABLE", &schema, Some(&table)).await?); - } - - tables.sort(); - let mut table_ddls = Vec::with_capacity(tables.len()); - for table in tables { - table_ddls.push(self.show_create("TABLE", &schema, Some(&table)).await?); - } - - views.sort(); - let mut view_ddls = Vec::with_capacity(views.len()); - for view in views { - view_ddls.push(self.show_create("VIEW", &schema, Some(&view)).await?); - } - - let ddl = build_schema_ddl( - &schema, - create_database, - physical_ddls, - table_ddls, - view_ddls, - ); - ddl_by_schema.push((schema, ddl)); - } - - Ok(ddl_by_schema) + async fn build_schema_ddl(&self, schema: &str) -> Result { + let (mut physical_tables, mut tables, mut views) = self.get_schema_objects(schema).await?; + physical_tables.sort(); + tables.sort(); + views.sort(); + let objects = std::iter::once(("DATABASE", None)) + .chain( + physical_tables + .iter() + .chain(&tables) + .map(|table| ("TABLE", Some(table.as_str()))), + ) + .chain(views.iter().map(|view| ("VIEW", Some(view.as_str())))); + let mut ddl = format!("-- Schema: {schema}\n"); + append_schema_ddl( + &self.database_client, + &self.config.catalog, + schema, + objects, + &mut ddl, + ) + .await + .context(DatabaseSnafu)?; + ddl.push('\n'); + Ok(ddl) } async fn get_schema_objects( @@ -734,68 +759,17 @@ impl ExportCreate { Ok(tables.into_iter().collect()) } - - async fn show_create( - &self, - show_type: &str, - schema: &str, - table: Option<&str>, - ) -> Result { - let sql = match table { - Some(table) => format!( - r#"SHOW CREATE {} "{}"."{}"."{}""#, - show_type, - escape_sql_identifier(&self.config.catalog), - escape_sql_identifier(schema), - escape_sql_identifier(table) - ), - None => format!( - r#"SHOW CREATE {} "{}"."{}""#, - show_type, - escape_sql_identifier(&self.config.catalog), - escape_sql_identifier(schema) - ), - }; - - let records: Option>> = self - .database_client - .sql_in_public(&sql) - .await - .context(DatabaseSnafu)?; - let rows = records.context(EmptyResultSnafu)?; - let row = rows.first().context(EmptyResultSnafu)?; - let Some(Value::String(create)) = row.get(1) else { - return UnexpectedValueTypeSnafu.fail(); - }; - - Ok(format!("{};\n", create)) - } -} - -fn build_schema_ddl( - schema: &str, - create_database: String, - physical_tables: Vec, - tables: Vec, - views: Vec, -) -> String { - let mut ddl = String::new(); - ddl.push_str(&format!("-- Schema: {}\n", schema)); - ddl.push_str(&create_database); - for stmt in physical_tables { - ddl.push_str(&stmt); - } - for stmt in tables { - ddl.push_str(&stmt); - } - for stmt in views { - ddl.push_str(&stmt); - } - ddl.push('\n'); - ddl } fn validate_resume_config(manifest: &Manifest, config: &ExportConfig) -> Result<()> { + if config.packed && !manifest.schema_only && !manifest.is_packed() { + return ResumeConfigMismatchSnafu { + field: "metric_data_layout", + existing: "standalone".to_string(), + requested: "packed".to_string(), + } + .fail(); + } if manifest.schema_only != config.schema_only { return SchemaOnlyModeMismatchSnafu { existing_schema_only: manifest.schema_only, @@ -1593,25 +1567,6 @@ mod tests { ); } - #[test] - fn test_build_schema_ddl_order() { - let ddl = build_schema_ddl( - "public", - "CREATE DATABASE public;\n".to_string(), - vec!["PHYSICAL;\n".to_string()], - vec!["TABLE;\n".to_string()], - vec!["VIEW;\n".to_string()], - ); - - let db_pos = ddl.find("CREATE DATABASE").unwrap(); - let physical_pos = ddl.find("PHYSICAL;").unwrap(); - let table_pos = ddl.find("TABLE;").unwrap(); - let view_pos = ddl.find("VIEW;").unwrap(); - assert!(db_pos < physical_pos); - assert!(physical_pos < table_pos); - assert!(table_pos < view_pos); - } - #[tokio::test] async fn test_build_rejects_chunk_window_without_bounds() { let cmd = ExportCreateCommand::parse_from([ @@ -1801,6 +1756,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1820,7 +1776,7 @@ mod tests { #[test] fn test_validate_resume_config_accepts_schema_selection_with_different_case_and_order() { - let manifest = Manifest::new_for_export( + let mut manifest = Manifest::new_for_export( "greptime".to_string(), vec!["public".to_string(), "analytics".to_string()], false, @@ -1829,7 +1785,7 @@ mod tests { None, ) .unwrap(); - let config = ExportConfig { + let mut config = ExportConfig { catalog: "greptime".to_string(), schemas: Some(vec![ "ANALYTICS".to_string(), @@ -1839,6 +1795,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1850,6 +1807,18 @@ mod tests { }; assert!(validate_resume_config(&manifest, &config).is_ok()); + config.packed = true; + assert!( + validate_resume_config(&manifest, &config) + .unwrap_err() + .to_string() + .contains("metric_data_layout") + ); + manifest.version = 2; + manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); + assert!(validate_resume_config(&manifest, &config).is_ok()); + config.packed = false; + assert!(validate_resume_config(&manifest, &config).is_ok()); } #[test] @@ -1872,6 +1841,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range, chunk_time_window: Some(Duration::from_secs(3600)), @@ -1906,6 +1876,7 @@ mod tests { schema_only: false, format: DataFormat::Csv, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1942,6 +1913,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::new(Some(start), Some(start)), chunk_time_window: None, diff --git a/src/cli/src/data/export_v2/coordinator.rs b/src/cli/src/data/export_v2/coordinator.rs index 5f95c43603e..c6181e51211 100644 --- a/src/cli/src/data/export_v2/coordinator.rs +++ b/src/cli/src/data/export_v2/coordinator.rs @@ -40,6 +40,7 @@ struct ExportContext<'a> { format: DataFormat, parallelism: usize, experimental_metric_export: bool, + packed: bool, resume: bool, } @@ -71,6 +72,7 @@ pub async fn export_data( catalog: manifest.catalog.clone(), schemas: manifest.schemas.clone(), format: manifest.format, + packed: manifest.is_packed(), parallelism: options.parallelism, experimental_metric_export: options.experimental_metric_export, resume: options.resume, @@ -279,7 +281,7 @@ async fn export_chunk( if context.experimental_metric_export { context .storage - .prepare_export_chunk(&context.schemas, chunk_id, context.resume) + .prepare_export_chunk(&context.schemas, chunk_id, context.resume, context.packed) .await?; } let scheme = StorageScheme::from_uri(context.snapshot_uri)?; @@ -289,6 +291,7 @@ async fn export_chunk( time_range, parallelism: context.parallelism, experimental_metric_export: context.experimental_metric_export, + packed: context.packed, }; for schema in &context.schemas { @@ -314,6 +317,39 @@ async fn export_chunk( } let files = list_chunk_files(context.storage, &context.schemas, chunk_id).await?; + if context.packed { + use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; + let mut expected = Vec::new(); + for schema in &context.schemas { + let prefix = data_dir_for_schema_chunk(schema, chunk_id); + let index_path = format!("{prefix}{PACK_INDEX_FILE}"); + let text = context.storage.read_text(&index_path).await?; + let index: PackIndex = serde_json::from_str(&text).map_err(|e| { + crate::data::export_v2::error::InvalidUriSnafu { + uri: context.snapshot_uri, + reason: e.to_string(), + } + .build() + })?; + index.validate().map_err(|e| { + crate::data::export_v2::error::InvalidUriSnafu { + uri: context.snapshot_uri, + reason: e.to_string(), + } + .build() + })?; + expected.push(index_path); + expected.extend(index.objects.iter().map(|o| format!("{prefix}{}", o.path))); + } + expected.sort(); + if files != expected { + return crate::data::export_v2::error::InvalidUriSnafu { + uri: context.snapshot_uri, + reason: "chunk inventory differs from packed indexes", + } + .fail(); + } + } info!("Collected {} files for chunk {}", files.len(), chunk_id); Ok(files) } @@ -462,6 +498,158 @@ mod tests { } } + struct PausedClose { + inner: Option, + started: std::sync::Arc, + release: std::sync::Arc, + } + + impl object_store::layers::mock::oio::Write for PausedClose { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.inner.as_mut().unwrap().write(bytes).await + } + + async fn close(&mut self) -> object_store::Result { + let mut inner = self.inner.take().unwrap(); + let (started, release) = (self.started.clone(), self.release.clone()); + // Model blocking storage I/O that survives dropping its caller. + let (inner, result) = tokio::spawn(async move { + started.notify_one(); + release.notified().await; + let result = inner.close().await; + (inner, result) + }) + .await + .unwrap(); + self.inner = Some(inner); + result + } + + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.as_mut().unwrap().abort().await + } + } + + #[tokio::test] + async fn packed_cancellation_drains_close_before_failing_chunk() { + use std::sync::Arc; + + use common_datasource::packed_snapshot::PACK_INDEX_FILE; + use common_datasource::packed_writer::{PackedTableWriter, PackedWriter}; + use common_datasource::parquet_writer::ParquetFileWriter; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + use tokio_util::sync::CancellationToken; + + use crate::data::snapshot_storage::OpenDalStorage; + + for paused_path in ["pack-000000.bin", PACK_INDEX_FILE] { + let directory = tempfile::tempdir().unwrap(); + let uri = url::Url::from_directory_path(directory.path()).unwrap(); + let storage = + OpenDalStorage::from_uri(uri.as_str(), &ObjectStoreConfig::default()).unwrap(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let factory: MockWriterFactory = Arc::new({ + let (started, release) = (started.clone(), release.clone()); + move |path, _, inner| { + if path.trim_start_matches('/') == paused_path { + Box::new(PausedClose { + inner: Some(inner), + started: started.clone(), + release: release.clone(), + }) + } else { + inner + } + } + }); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let packed = PackedWriter::new(store.clone()).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + true, + )])); + let mut table = ParquetFileWriter::open_packed( + schema, + store.clone(), + "unused", + None, + PackedTableWriter::new(packed.clone(), "empty".into(), 1, false), + ) + .unwrap(); + table.finish(None).await.unwrap(); + let token = CancellationToken::new(); + let mut manifest = pending_manifest(1); + manifest.version = 2; + manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); + let export = export_data_concurrent( + &storage, + &mut manifest, + 2, + &crate::data::progress::NoopProgress, + |_, _| async { + let mut packed = packed.lock().await; + let result = packed.finish(&token).await; + if result.is_err() { + packed.abort().await.unwrap(); + } + result.map_err(|error| { + crate::data::export_v2::error::IoSnafu { + operation: "exporting packed chunk", + error: std::io::Error::other(error), + } + .build() + }) + }, + ); + tokio::pin!(export); + tokio::select! { + _ = started.notified() => {}, + result = &mut export => panic!("completed while close paused: {result:?}"), + } + token.cancel(); + assert!(futures::poll!(&mut export).is_pending()); + assert_eq!( + storage.read_manifest().await.unwrap().chunks[0].status, + ChunkStatus::InProgress + ); + release.notify_one(); + assert!(export.await.unwrap_err().to_string().contains("cancelled")); + let persisted = storage.read_manifest().await.unwrap(); + assert_eq!(persisted.chunks[0].status, ChunkStatus::Failed); + assert!(!persisted.is_complete()); + assert!(persisted.chunks[0].files.is_empty()); + assert!( + store + .stat("pack-000000.bin") + .await + .unwrap() + .content_length() + > 0 + ); + assert_eq!( + store.exists(PACK_INDEX_FILE).await.unwrap(), + paused_path == PACK_INDEX_FILE + ); + if paused_path == PACK_INDEX_FILE { + let index: common_datasource::packed_snapshot::PackIndex = + serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()) + .unwrap(); + index.validate_membership(["empty"]).unwrap(); + } + } + } + #[test] fn test_next_eligible_chunk_scans_in_order() { let manifest = pending_manifest(3); diff --git a/src/cli/src/data/export_v2/data.rs b/src/cli/src/data/export_v2/data.rs index 8723154e4b1..53c3899f59e 100644 --- a/src/cli/src/data/export_v2/data.rs +++ b/src/cli/src/data/export_v2/data.rs @@ -31,6 +31,7 @@ pub(super) struct CopyOptions { pub(super) time_range: TimeRange, pub(super) parallelism: usize, pub(crate) experimental_metric_export: bool, + pub(super) packed: bool, } pub(super) struct CopyTarget { @@ -220,6 +221,9 @@ fn build_with_options(options: &CopyOptions) -> String { if options.experimental_metric_export { parts.push("experimental_metric_export='true'".to_string()); } + if options.packed { + parts.push("metric_data_layout='packed'".to_string()); + } if let Some(start) = options.time_range.start { parts.push(format!( "START_TIME='{}'", diff --git a/src/cli/src/data/export_v2/manifest.rs b/src/cli/src/data/export_v2/manifest.rs index 5dd9d6f1684..d56f2732490 100644 --- a/src/cli/src/data/export_v2/manifest.rs +++ b/src/cli/src/data/export_v2/manifest.rs @@ -28,7 +28,7 @@ use crate::data::export_v2::error::{ TimeParseInvalidFormatSnafu, }; -/// Manifest format version produced by the current exporter. +/// Default standalone manifest version. Packed data snapshots use version 2. pub const MANIFEST_VERSION: u32 = 1; /// Manifest file name within snapshot directory. diff --git a/src/cli/src/data/import_v2/command.rs b/src/cli/src/data/import_v2/command.rs index 33cf9cb1ffb..62b9de08816 100644 --- a/src/cli/src/data/import_v2/command.rs +++ b/src/cli/src/data/import_v2/command.rs @@ -222,12 +222,19 @@ impl Import { manifest.validate_layout().map_err(|reason| { crate::data::import_v2::error::InvalidPackedSnapshotSnafu { reason }.build() })?; - if manifest.is_packed() { - self.database_client - .require_packed_import() + let capabilities = if manifest.is_packed() { + let capabilities = self + .database_client + .capabilities() .await .context(crate::data::import_v2::error::DatabaseSnafu)?; - } + capabilities + .require("metric_packed_import") + .context(crate::data::import_v2::error::DatabaseSnafu)?; + Some(capabilities) + } else { + None + }; info!("Snapshot contains {} schema(s)", manifest.schemas.len()); @@ -279,6 +286,15 @@ impl Import { return Ok(()); } + let capabilities = match capabilities { + Some(capabilities) => capabilities, + None => self + .database_client + .capabilities() + .await + .context(crate::data::import_v2::error::DatabaseSnafu)?, + }; + let mut resume_session = if !data_tasks.is_empty() { let state_path = resolve_state_path( self.state_path.as_deref(), @@ -316,7 +332,9 @@ impl Import { false } else { let executor = DdlExecutor::new(&self.database_client); - executor.execute_strict(&ddl_statements).await?; + executor + .execute_strict(&ddl_statements, capabilities.supports("metric_batch_ddl")) + .await?; if let Some(session) = resume_session.as_mut() { session.mark_ddl_completed().await?; } @@ -356,9 +374,7 @@ impl Import { .await .context(SnapshotStorageSnafu)?; statements.extend( - parse_ddl_statements(&content) - .into_iter() - .map(|sql| ddl_statement_for_schema(schema, sql)), + iter_ddl_statements(&content).map(|sql| ddl_statement_for_schema(schema, sql)), ); } @@ -402,90 +418,92 @@ impl ImportTaskExecutor for CopyDatabaseImportTaskExecutor<'_> { } } +#[cfg(test)] fn parse_ddl_statements(content: &str) -> Vec { - let mut statements = Vec::new(); - let mut current = String::new(); + iter_ddl_statements(content).collect() +} + +pub(crate) fn iter_ddl_statements(content: &str) -> impl Iterator + '_ { let mut chars = content.chars().peekable(); let mut in_single_quote = false; let mut in_double_quote = false; let mut in_line_comment = false; let mut in_block_comment = false; - while let Some(ch) = chars.next() { - if in_line_comment { - if ch == '\n' { - in_line_comment = false; - current.push('\n'); - } - continue; - } - - if in_block_comment { - if ch == '*' && chars.peek() == Some(&'/') { - chars.next(); - in_block_comment = false; - } - continue; - } - - if in_single_quote { - current.push(ch); - if ch == '\'' { - if chars.peek() == Some(&'\'') { - current.push(chars.next().expect("peeked quote must exist")); - } else { - in_single_quote = false; + std::iter::from_fn(move || { + let mut current = String::new(); + while let Some(ch) = chars.next() { + if in_line_comment { + if ch == '\n' { + in_line_comment = false; + current.push('\n'); } + continue; } - continue; - } - if in_double_quote { - current.push(ch); - if ch == '"' { - if chars.peek() == Some(&'"') { - current.push(chars.next().expect("peeked quote must exist")); - } else { - in_double_quote = false; + if in_block_comment { + if ch == '*' && chars.peek() == Some(&'/') { + chars.next(); + in_block_comment = false; } + continue; } - continue; - } - match ch { - '-' if chars.peek() == Some(&'-') => { - chars.next(); - in_line_comment = true; - } - '/' if chars.peek() == Some(&'*') => { - chars.next(); - in_block_comment = true; - } - '\'' => { - in_single_quote = true; + if in_single_quote { current.push(ch); - } - '"' => { - in_double_quote = true; - current.push(ch); - } - ';' => { - let statement = current.trim(); - if !statement.is_empty() { - statements.push(statement.to_string()); + if ch == '\'' { + if chars.peek() == Some(&'\'') { + current.push(chars.next().expect("peeked quote must exist")); + } else { + in_single_quote = false; + } } - current.clear(); + continue; + } + + if in_double_quote { + current.push(ch); + if ch == '"' { + if chars.peek() == Some(&'"') { + current.push(chars.next().expect("peeked quote must exist")); + } else { + in_double_quote = false; + } + } + continue; + } + + match ch { + '-' if chars.peek() == Some(&'-') => { + chars.next(); + in_line_comment = true; + } + '/' if chars.peek() == Some(&'*') => { + chars.next(); + in_block_comment = true; + } + '\'' => { + in_single_quote = true; + current.push(ch); + } + '"' => { + in_double_quote = true; + current.push(ch); + } + ';' => { + let statement = current.trim(); + if !statement.is_empty() { + return Some(statement.to_string()); + } + current.clear(); + } + _ => current.push(ch), } - _ => current.push(ch), } - } - let statement = current.trim(); - if !statement.is_empty() { - statements.push(statement.to_string()); - } - - statements + let statement = current.trim(); + (!statement.is_empty()).then(|| statement.to_string()) + }) } fn ddl_statement_for_schema(schema: &str, sql: String) -> DdlStatement { diff --git a/src/cli/src/data/import_v2/executor.rs b/src/cli/src/data/import_v2/executor.rs index 3f2bf66ae68..6f152bc5c6d 100644 --- a/src/cli/src/data/import_v2/executor.rs +++ b/src/cli/src/data/import_v2/executor.rs @@ -14,8 +14,18 @@ //! DDL execution for import. +use std::collections::HashSet; + +use common_catalog::consts::DEFAULT_SCHEMA_NAME; use common_telemetry::info; +use servers::query_handler::sql::{MAX_LOGICAL_TABLE_DDL_BYTES, MAX_LOGICAL_TABLE_DDL_STATEMENTS}; use snafu::ResultExt; +use sql::dialect::GreptimeDbDialect; +use sql::parser::{ParseOptions, ParserContext}; +use sql::statements::statement::Statement; +use store_api::metric_engine_consts::{ + LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, +}; use crate::data::import_v2::error::{DatabaseSnafu, Result}; use crate::database::DatabaseClient; @@ -54,69 +64,379 @@ impl<'a> DdlExecutor<'a> { Self { client } } - /// Executes a list of DDL statements, stopping on first error. - pub async fn execute_strict(&self, statements: &[DdlStatement]) -> Result<()> { - let total = statements.len(); - + /// Executes DDL in order, stopping at the first failed request. + pub async fn execute_strict( + &self, + statements: &[DdlStatement], + batch_supported: bool, + ) -> Result<()> { + let mut batch = LogicalBatch::default(); for (i, stmt) in statements.iter().enumerate() { - let preview = preview_sql(&stmt.sql); - - info!("Executing DDL ({}/{}): {}", i + 1, total, preview); - - if let Some(schema) = stmt.execution_schema.as_deref() { + info!("Executing DDL ({}/{})", i + 1, statements.len()); + let schema = stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME); + let candidate = if batch_supported { + classify(stmt, self.client.catalog()).context(DatabaseSnafu)? + } else { + None + }; + if let Some(candidate) = candidate { + if !batch.accepts(&candidate, &stmt.sql) { + batch.flush(self.client).await?; + } + batch.push(candidate, &stmt.sql); + } else { + batch.flush(self.client).await?; self.client .sql(&stmt.sql, schema) .await .context(DatabaseSnafu)?; - } else { - self.client - .sql_in_public(&stmt.sql) - .await - .context(DatabaseSnafu)?; } } + batch.flush(self.client).await + } +} +struct LogicalTable { + catalog: String, + execution_schema: String, + schema: String, + physical: String, + name: String, +} + +fn classify(stmt: &DdlStatement, catalog: &str) -> crate::error::Result> { + if stmt.sql.len() > MAX_LOGICAL_TABLE_DDL_BYTES { + return Ok(None); + } + let parsed = parse_ddl(&stmt.sql)?; + let [Statement::CreateTable(create)] = parsed.as_slice() else { + return Ok(None); + }; + let Some(physical) = create.options.get(LOGICAL_TABLE_METADATA_KEY) else { + return Ok(None); + }; + if create.engine != METRIC_ENGINE_NAME + || create.options.value(PHYSICAL_TABLE_METADATA_KEY).is_some() + { + return Ok(None); + } + let parts: Vec<_> = create + .name + .0 + .iter() + .filter_map(|p| p.as_ident().map(|i| i.value.as_str())) + .collect(); + let schema = stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME); + let (catalog, schema, name) = match parts.as_slice() { + [name] => (catalog, schema, *name), + [schema, name] => (catalog, *schema, *name), + [catalog, schema, name] => (*catalog, *schema, *name), + _ => return Ok(None), + }; + Ok(Some(LogicalTable { + catalog: catalog.into(), + execution_schema: stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME) + .into(), + schema: schema.into(), + physical: physical.into(), + name: name.into(), + })) +} + +pub(crate) fn parse_ddl(sql: &str) -> crate::error::Result> { + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()).map_err( + |_| { + crate::error::InvalidArgumentsSnafu { + msg: "cannot parse import DDL", + } + .build() + }, + ) +} + +#[derive(Default)] +struct LogicalBatch { + group: Option, + names: HashSet, + sql: String, +} + +impl LogicalBatch { + fn accepts(&self, table: &LogicalTable, sql: &str) -> bool { + self.group.as_ref().is_none_or(|g| { + g.catalog == table.catalog + && g.schema == table.schema + && g.physical == table.physical + && g.execution_schema == table.execution_schema + }) && !self.names.contains(&table.name) + && self.names.len() < MAX_LOGICAL_TABLE_DDL_STATEMENTS + && self.sql.len() + (if self.sql.is_empty() { 0 } else { 2 }) + sql.len() + <= MAX_LOGICAL_TABLE_DDL_BYTES + } + + fn push(&mut self, table: LogicalTable, sql: &str) { + if !self.sql.is_empty() { + self.sql.push_str(";\n"); + } + self.sql.push_str(sql); + self.names.insert(table.name.clone()); + self.group = Some(table); + } + + async fn flush(&mut self, client: &DatabaseClient) -> Result<()> { + if let Some(group) = self.group.as_ref() { + client + .logical_tables(&self.sql, &group.execution_schema, self.names.len()) + .await + .context(DatabaseSnafu)?; + self.sql.clear(); + self.names.clear(); + self.group = None; + } Ok(()) } } -fn preview_sql(sql: &str) -> String { - let mut chars = sql.chars(); - let preview: String = chars.by_ref().take(80).collect(); - if chars.next().is_some() { - format!("{preview}...") - } else { - preview - } -} #[cfg(test)] mod tests { use super::*; - #[test] - fn test_statement_without_execution_schema_uses_public() { - let stmt = DdlStatement::new("CREATE DATABASE IF NOT EXISTS test_db".to_string()); - assert_eq!(stmt.execution_schema, None); + fn logical(name: &str, physical: &str) -> String { + format!( + "CREATE TABLE IF NOT EXISTS {name} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='{physical}')" + ) } #[test] - fn test_statement_with_execution_schema_preserves_context() { - let stmt = DdlStatement::with_execution_schema( - r#"CREATE TABLE IF NOT EXISTS "my""schema"."metrics" (ts TIMESTAMP TIME INDEX)"# - .to_string(), - r#"my"schema"#.to_string(), - ); - assert_eq!(stmt.execution_schema.as_deref(), Some(r#"my"schema"#)); + fn classification_and_batch_boundaries() { + let mut batch = LogicalBatch::default(); + for (name, context, physical, expected) in [ + ("A", "public", "p.q", true), + ("public.b", "other", "p.q", false), + ("greptime.public.c", "public", "p.q", true), + ("a", "public", "p.q", false), + ("\"A\"", "public", "p.q", true), + ("d", "public", "p", false), + ("d", "other", "p.q", false), + ("other.public.d", "public", "p.q", false), + ] { + let sql = logical(name, physical); + let stmt = DdlStatement::with_execution_schema(sql.clone(), context.into()); + let table = classify(&stmt, "greptime").unwrap().unwrap(); + assert_eq!(batch.accepts(&table, &sql), expected, "{name}"); + if expected { + batch.push(table, &sql); + } + } + for sql in [ + "CREATE DATABASE d", + "CREATE TABLE t (ts TIMESTAMP TIME INDEX)", + "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table='')", + "CREATE VIEW v AS SELECT * FROM a", + ] { + assert!( + classify(&DdlStatement::new(sql.into()), "greptime") + .unwrap() + .is_none() + ); + } + let mut batch = LogicalBatch::default(); + for i in 0..128 { + let sql = logical(&format!("t{i}"), "p"); + let table = classify(&DdlStatement::new(sql.clone()), "greptime") + .unwrap() + .unwrap(); + assert!(batch.accepts(&table, &sql)); + batch.push(table, &sql); + } + let table = classify(&DdlStatement::new(logical("next", "p")), "greptime") + .unwrap() + .unwrap(); + assert!(!batch.accepts(&table, &logical("next", "p"))); + batch.names.clear(); + batch.sql = "测".repeat(MAX_LOGICAL_TABLE_DDL_BYTES / 3 - 1); + assert!(batch.accepts(&table, "ab")); + assert!(!batch.accepts(&table, "测")); + let oversized = + DdlStatement::new(logical("large", &"x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES))); + assert!(classify(&oversized, "greptime").unwrap().is_none()); } - #[test] - fn test_preview_sql_truncates_at_char_boundary() { - let sql = format!( - "CREATE TABLE {} (ts TIMESTAMP TIME INDEX)", - "测".repeat(100) - ); - let preview = preview_sql(&sql); - assert!(preview.ends_with("...")); - assert!(preview.is_char_boundary(preview.len())); + #[tokio::test] + async fn oversized_ddl_uses_ordinary_sql_in_order() { + let large = "x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES); + let statements = [ + DdlStatement::new(logical("before", "p")), + DdlStatement::new(format!( + "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, payload STRING DEFAULT '{large}')" + )), + DdlStatement::new(logical("between", "p")), + DdlStatement::new(format!("CREATE VIEW v AS SELECT '{large}' AS payload")), + DdlStatement::new(logical("oversized", &large)), + DdlStatement::new(logical("after", "p")), + ]; + let (client, requests, server) = crate::database::tests::test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await + .unwrap(); + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), statements.len()); + for (i, (request, statement)) in requests.iter().zip(&statements).enumerate() { + let endpoint = if [0, 2, 5].contains(&i) { + "ddl/logical-tables" + } else { + "sql" + }; + assert!(request.starts_with(&format!("POST /v1/{endpoint} "))); + let body = request.split_once("\r\n\r\n").unwrap().1; + let sql = url::form_urlencoded::parse(body.as_bytes()) + .find(|(key, _)| key == "sql") + .unwrap() + .1; + assert_eq!(sql, statement.sql); + } + } + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn qualified_targets_preserve_request_context() { + let (client, requests, server) = crate::database::tests::test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + let statements = [ + DdlStatement::new(logical("other.metrics.a", "p")), + DdlStatement::with_execution_schema( + logical("other.metrics.b", "p"), + "existing_schema".into(), + ), + ]; + DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await + .unwrap(); + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + for (request, schema) in requests.iter().zip(["public", "existing_schema"]) { + assert!(request.starts_with("POST /v1/ddl/logical-tables ")); + let body = request.split_once("\r\n\r\n").unwrap().1; + let db = url::form_urlencoded::parse(body.as_bytes()) + .find(|(k, _)| k == "db") + .unwrap() + .1; + assert_eq!(db, format!("greptime-{schema}")); + } + } + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn transport_preserves_order_and_stops_without_fallback() { + use crate::database::tests::test_server; + let statements = [ + DdlStatement::new("CREATE DATABASE d".into()), + DdlStatement::new(logical("a", "p").replace( + "on_physical_table='p'", + "on_physical_table='p', access_key_id='original-secret'", + )), + DdlStatement::new(logical("a", "p")), + DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()), + ]; + for supported in [true, false] { + let (client, requests, server) = test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + DdlExecutor::new(&client) + .execute_strict(&statements, supported) + .await + .unwrap(); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 4); + for (i, request) in requests.iter().enumerate() { + let path = if supported && (i == 1 || i == 2) { + "ddl/logical-tables" + } else { + "sql" + }; + assert!(request.starts_with(&format!("POST /v1/{path} "))); + assert!( + request + .to_lowercase() + .contains("authorization: basic dxnlcjpwyxnzd29yza==") + ); + let body = request.split_once("\r\n\r\n").unwrap().1; + let form: std::collections::HashMap<_, _> = + url::form_urlencoded::parse(body.as_bytes()) + .into_owned() + .collect(); + assert_eq!(form["sql"], statements[i].sql); + assert_eq!(form["db"], "greptime-public"); + } + server.abort(); + } + let statements = [ + DdlStatement::new(logical("a", "p")), + DdlStatement::new(logical("b", "p")), + DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()), + ]; + for (status, body, success) in [ + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":0}]}"#, + true, + ), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + false, + ), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":1}]}"#, + false, + ), + (200, r#"{"execution_time_ms":0,"output":[{},{}]}"#, false), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"error":"secret"}]}"#, + false, + ), + (500, r#"{"error":"secret"}"#, false), + (200, "invalid", false), + ] { + let (client, requests, server) = test_server(status, body).await; + let result = DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await; + assert_eq!(result.is_ok(), success); + assert_eq!(requests.lock().unwrap().len(), if success { 2 } else { 1 }); + if let Err(error) = result { + assert!(!format!("{error:?}").contains("secret")); + } + server.abort(); + let _ = server.await; + } } } diff --git a/src/cli/src/data/import_v2/packed.rs b/src/cli/src/data/import_v2/packed.rs index 303c866ba69..c2f1ec35889 100644 --- a/src/cli/src/data/import_v2/packed.rs +++ b/src/cli/src/data/import_v2/packed.rs @@ -16,12 +16,11 @@ use std::collections::HashSet; use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; use snafu::ResultExt; -use sql::dialect::GreptimeDbDialect; -use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use crate::data::export_v2::manifest::{ChunkStatus, Manifest}; use crate::data::import_v2::error::{InvalidPackedSnapshotSnafu, Result, SnapshotStorageSnafu}; +use crate::data::import_v2::{command, executor}; use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; use crate::data::snapshot_storage::SnapshotStorage; @@ -51,49 +50,46 @@ pub(crate) async fn validate_snapshot( .read_text(&ddl_path_for_schema(schema)) .await .context(SnapshotStorageSnafu)?; - let statements = ParserContext::create_with_dialect( - &ddl, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .map_err(|e| { - InvalidPackedSnapshotSnafu { - reason: e.to_string(), - } - .build() - })?; let mut names = HashSet::new(); - for statement in statements { - if let Statement::CreateTable(create) = statement { - if create.engine.eq_ignore_ascii_case("metric") - && create.options.get("on_physical_table").is_none() - { - continue; + for sql in command::iter_ddl_statements(&ddl) { + let statements = executor::parse_ddl(&sql).map_err(|_| { + InvalidPackedSnapshotSnafu { + reason: "invalid snapshot DDL", } - let parts: Vec<_> = create - .name - .0 - .iter() - .map(|p| p.as_ident().map(|i| i.value.as_str())) - .collect(); - let allowed = match parts.as_slice() { - [Some(_)] => true, - [Some(s), Some(_)] => *s == schema, - [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema, - _ => false, - }; - if !allowed { - return InvalidPackedSnapshotSnafu { - reason: "snapshot table DDL escapes its catalog/schema", + .build() + })?; + for statement in statements { + if let Statement::CreateTable(create) = statement { + if create.engine.eq_ignore_ascii_case("metric") + && create.options.get("on_physical_table").is_none() + { + continue; } - .fail(); - } - let name = parts.last().and_then(|p| *p).unwrap_or_default(); - if !names.insert(name.to_string()) { - return InvalidPackedSnapshotSnafu { - reason: "duplicate snapshot table DDL", + let parts: Vec<_> = create + .name + .0 + .iter() + .map(|p| p.as_ident().map(|i| i.value.as_str())) + .collect(); + let allowed = match parts.as_slice() { + [Some(_)] => true, + [Some(s), Some(_)] => *s == schema, + [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema, + _ => false, + }; + if !allowed { + return InvalidPackedSnapshotSnafu { + reason: "snapshot table DDL escapes its catalog/schema", + } + .fail(); + } + let name = parts.last().and_then(|p| *p).unwrap_or_default(); + if !names.insert(name.to_string()) { + return InvalidPackedSnapshotSnafu { + reason: "duplicate snapshot table DDL", + } + .fail(); } - .fail(); } } } @@ -180,13 +176,14 @@ mod tests { .unwrap(); let uri = url::Url::from_file_path(dir.path()).unwrap(); let storage = OpenDalStorage::from_uri(uri.as_str(), &Default::default()).unwrap(); - storage.write_text("schema/ddl/public.sql", "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\";").await.unwrap(); + let large = "x".repeat(servers::query_handler::sql::MAX_LOGICAL_TABLE_DDL_BYTES); + storage.write_text("schema/ddl/public.sql", &format!("CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\"; CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, payload STRING DEFAULT '{large}'); CREATE VIEW large_view AS SELECT '{large}' AS payload;")).await.unwrap(); let path = "data/public/1/pack-index.json"; - let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":12}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0}]}); + let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":24}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0},{"table_name":"ordinary","object":"pack-0.bin","offset":12,"length":12,"row_count":0}]}); storage.write_text(path, &index.to_string()).await.unwrap(); assert!(dir.path().join(path).is_file()); storage - .write_text("data/public/1/pack-0.bin", "abcdefghijkl") + .write_text("data/public/1/pack-0.bin", "abcdefghijklmnopqrstuvwx") .await .unwrap(); let mut manifest = Manifest::new_schema_only("greptime".into(), vec!["public".into()]); diff --git a/src/cli/src/data/schema_export.rs b/src/cli/src/data/schema_export.rs new file mode 100644 index 00000000000..8bb99121d12 --- /dev/null +++ b/src/cli/src/data/schema_export.rs @@ -0,0 +1,508 @@ +// 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. + +//! Bounded SHOW CREATE requests shared by both export commands. + +use common_catalog::consts::DEFAULT_SCHEMA_NAME; +use servers::http::GreptimeQueryOutput; +use servers::http::result::greptime_result_v1::GreptimedbV1Response; + +use crate::data::sql::escape_sql_identifier; +use crate::database::DatabaseClient; +use crate::error::{InvalidArgumentsSnafu, Result, UnexpectedSnafu}; + +const MAX_STATEMENTS: usize = 128; +const MAX_SQL_BYTES: usize = 128 * 1024; + +/// Appends DDL in dependency order supplied by the caller using bounded requests. +pub(crate) fn append_schema_ddl<'a>( + client: &'a DatabaseClient, + catalog: &'a str, + schema: &'a str, + objects: impl Iterator)> + Send + 'a, + ddl: &'a mut String, +) -> impl std::future::Future> + Send + 'a { + let mut statements = objects + .map(move |(kind, table)| show_create_sql(kind, catalog, schema, table)) + .peekable(); + async move { + loop { + let batch = next_batch(&mut statements)?; + if batch.is_empty() { + return Ok(()); + } + let response = client + .sql_response(&batch.concat(), DEFAULT_SCHEMA_NAME) + .await + .map_err(|error| { + UnexpectedSnafu { + msg: format!("SHOW CREATE batch failed for {}: {error}", batch.concat()), + } + .build() + })?; + append_results(&batch, &response, ddl)?; + } + } +} + +fn show_create_sql(kind: &str, catalog: &str, schema: &str, table: Option<&str>) -> String { + let mut sql = format!( + "SHOW CREATE {kind} \"{}\".\"{}\"", + escape_sql_identifier(catalog), + escape_sql_identifier(schema), + ); + if let Some(table) = table { + sql.push_str(".\""); + sql.push_str(&escape_sql_identifier(table)); + sql.push('"'); + } + sql.push_str(";\n"); + sql +} + +fn next_batch( + statements: &mut std::iter::Peekable>, +) -> Result> { + let mut batch = Vec::new(); + let mut bytes = 0; + while let Some(sql) = statements.peek() { + if sql.len() > MAX_SQL_BYTES { + return InvalidArgumentsSnafu { + msg: format!( + "SHOW CREATE exceeds {MAX_SQL_BYTES} SQL bytes ({}): {sql}", + sql.len() + ), + } + .fail(); + } + if batch.len() == MAX_STATEMENTS || bytes + sql.len() > MAX_SQL_BYTES { + break; + } + bytes += sql.len(); + if let Some(sql) = statements.next() { + batch.push(sql); + } + } + Ok(batch) +} + +fn append_results( + batch: &[String], + response: &GreptimedbV1Response, + ddl: &mut String, +) -> Result<()> { + let outputs = response.output(); + if outputs.len() != batch.len() { + return UnexpectedSnafu { + msg: format!( + "SHOW CREATE expected {} results, received {} for {}", + batch.len(), + outputs.len(), + batch.concat() + ), + } + .fail(); + } + for (sql, output) in batch.iter().zip(outputs) { + let GreptimeQueryOutput::Records(records) = output else { + return UnexpectedSnafu { + msg: format!("expected SHOW CREATE records for {sql}"), + } + .fail(); + }; + if records.num_cols() != 2 || records.rows().len() != 1 || records.rows()[0].len() != 2 { + return UnexpectedSnafu { + msg: format!("expected one two-column SHOW CREATE row for {sql}"), + } + .fail(); + } + let row = &records.rows()[0]; + let Some(create) = row[1].as_str().filter(|value| !value.is_empty()) else { + return UnexpectedSnafu { + msg: format!("expected SHOW CREATE DDL string for {sql}"), + } + .fail(); + }; + if !row[0].is_string() { + return UnexpectedSnafu { + msg: format!("expected SHOW CREATE object name for {sql}"), + } + .fail(); + } + ddl.push_str(create); + ddl.push_str(";\n"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; + + use super::*; + + fn records(name: &str, ddl: &str) -> Value { + json!({"records": {"schema": {"column_schemas": [ + {"name": "Table", "data_type": "String"}, + {"name": "Create Table", "data_type": "String"} + ]}, "rows": [[name, ddl]], "total_rows": 1}}) + } + + #[test] + fn batch_limits_include_quoted_utf8_and_separators() { + let special = show_create_sql("TABLE", "cat.name", "s\";库", Some("ta.\";表")); + assert_eq!( + special, + "SHOW CREATE TABLE \"cat.name\".\"s\"\";库\".\"ta.\"\";表\";\n" + ); + let mut statements = std::iter::repeat_n(special.clone(), 129).peekable(); + assert_eq!( + next_batch(&mut statements).unwrap(), + vec![special.clone(); 128] + ); + assert_eq!(next_batch(&mut statements).unwrap(), vec![special]); + assert!(next_batch(&mut statements).unwrap().is_empty()); + + let prefix = show_create_sql("TABLE", "c", "s", Some("")); + // Every quote expands to two bytes; this also exercises a multibyte name. + let name = format!("表{}", "\"".repeat((MAX_SQL_BYTES - prefix.len() - 3) / 2)); + let mut exact = show_create_sql("TABLE", "c", "s", Some(&name)); + let padding = MAX_SQL_BYTES - exact.len(); + exact = show_create_sql( + "TABLE", + "c", + "s", + Some(&format!("{name}{}", "a".repeat(padding))), + ); + assert_eq!(exact.len(), MAX_SQL_BYTES); + let small = show_create_sql("TABLE", "c", "s", Some("next")); + let mut statements = [exact.clone(), small.clone()].into_iter().peekable(); + assert_eq!(next_batch(&mut statements).unwrap(), vec![exact.clone()]); + assert_eq!(next_batch(&mut statements).unwrap(), vec![small]); + let mut oversized = std::iter::once(format!("{exact} ")).peekable(); + assert!( + next_batch(&mut oversized) + .unwrap_err() + .to_string() + .contains("exceeds 131072") + ); + + let half = show_create_sql( + "TABLE", + "c", + "s", + Some(&"a".repeat(MAX_SQL_BYTES / 2 - prefix.len())), + ); + let mut statements = [half.clone(), half.clone(), prefix.clone()] + .into_iter() + .peekable(); + assert_eq!( + next_batch(&mut statements).unwrap(), + vec![half.clone(), half] + ); + assert_eq!(next_batch(&mut statements).unwrap(), vec![prefix]); + } + + #[test] + fn show_create_requires_exact_record_shape_and_count() { + let batch = vec![ + "SHOW CREATE TABLE first".into(), + "SHOW CREATE TABLE middle".into(), + ]; + let valid = records("first", "CREATE TABLE first"); + for output in [ + vec![], + vec![valid.clone()], + vec![valid.clone(); 3], + vec![valid.clone(), json!({"affectedrows": 1})], + vec![valid.clone(), records("middle", "")], + vec![valid.clone(), { + let mut r = valid.clone(); + r["records"]["rows"] = json!([]); + r + }], + vec![valid.clone(), { + let mut r = valid.clone(); + r["records"]["rows"] = json!([["middle", 4]]); + r + }], + vec![valid.clone(), { + let mut r = valid.clone(); + r["records"]["rows"] = json!([["middle", "ddl", "extra"]]); + r + }], + ] { + let response = + serde_json::from_value(json!({"output": output, "execution_time_ms": 0})).unwrap(); + let error = append_results(&batch, &response, &mut String::new()).unwrap_err(); + assert!(error.to_string().contains("middle"), "{error}"); + } + } + + async fn start_mock_sql_server( + bodies: Vec<(u16, String)>, + ) -> (DatabaseClient, tokio::task::JoinHandle>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + tokio::time::timeout(std::time::Duration::from_secs(5), async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = BufReader::new(socket); + let mut queries = Vec::new(); + for (status, body) in bodies { + let mut headers = String::new(); + loop { + let mut line = String::new(); + if socket.read_line(&mut line).await.unwrap() == 0 { + assert!(headers.is_empty()); + return queries; + } + headers.push_str(&line); + if line == "\r\n" { + break; + } + } + let headers = headers.to_lowercase(); + assert!(headers.starts_with("post /v1/sql ")); + assert!(headers.contains("authorization: basic dxnlcjpwyxnzd29yza==")); + assert!(headers.contains("x-greptime-timeout: 7s")); + let length: usize = headers + .lines() + .find_map(|line| line.strip_prefix("content-length: ")) + .unwrap() + .parse() + .unwrap(); + let mut form = vec![0; length]; + socket.read_exact(&mut form).await.unwrap(); + let params: std::collections::HashMap<_, _> = + url::form_urlencoded::parse(&form).into_owned().collect(); + assert_eq!(params["db"], "catalog-public"); + queries.push(params["sql"].clone()); + socket + .get_mut() + .write_all( + format!( + "HTTP/1.1 {status} Response\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .await + .unwrap(); + } + queries + }) + .await + .expect("mock SQL exchange timed out") + }); + let client = DatabaseClient::new( + address.to_string(), + "catalog".into(), + Some("user:password".into()), + std::time::Duration::from_secs(7), + None, + true, + ); + (client, server) + } + + #[tokio::test] + async fn batches_preserve_order_context_and_reuse_connection() { + let names: Vec<_> = (0..129).map(|i| format!("表.{i}\";x")).collect(); + let outputs: Vec<_> = names + .iter() + .map(|name| records(name, &format!("DDL {name}"))) + .collect(); + let bodies = outputs + .chunks(128) + .map(|chunk| { + ( + 200, + json!({"output": chunk, "execution_time_ms": 0}).to_string(), + ) + }) + .collect(); + let (client, server) = start_mock_sql_server(bodies).await; + let mut ddl = String::new(); + append_schema_ddl( + &client.clone(), + "catalog", + "schema", + names.iter().map(|name| ("TABLE", Some(name.as_str()))), + &mut ddl, + ) + .await + .unwrap(); + assert_eq!( + ddl, + names + .iter() + .map(|name| format!("DDL {name};\n")) + .collect::() + ); + let queries = server.await.unwrap(); + let expected: Vec<_> = (0..129) + .map(|i| format!("SHOW CREATE TABLE \"catalog\".\"schema\".\"表.{i}\"\";x\";\n")) + .collect(); + assert_eq!( + queries, + vec![expected[..128].concat(), expected[128..].concat()] + ); + } + + #[tokio::test] + async fn legacy_command_stops_before_data_after_schema_failure() { + use clap::Parser; + + use crate::data::export::ExportCommand; + + let output = |rows: Value, columns: usize| { + json!({"records": {"schema": {"column_schemas": (0..columns) + .map(|i| json!({"name": i.to_string(), "data_type": "String"})) + .collect::>()}, "rows": rows, "total_rows": rows.as_array().unwrap().len()}}) + }; + let databases = output(json!([["public"]]), 1); + let mut malformed = records("middle", "CREATE TABLE middle"); + malformed["records"]["rows"][0][1] = json!(42); + let bodies = vec![ + vec![databases.clone()], + vec![records("public", "CREATE DATABASE public")], + vec![databases.clone()], + vec![output(json!([]), 3)], + vec![output( + json!([ + ["catalog", "public", "first", "BASE TABLE"], + ["catalog", "public", "middle", "BASE TABLE"], + ["catalog", "public", "last", "BASE TABLE"] + ]), + 4, + )], + vec![ + records("first", "CREATE TABLE first"), + malformed, + records("last", "CREATE TABLE last"), + ], + vec![databases], + vec![json!({"affectedrows": 0})], + ] + .into_iter() + .map(|output| { + ( + 200, + json!({"output": output, "execution_time_ms": 0}).to_string(), + ) + }) + .collect::>(); + for target in ["schema", "all"] { + let (client, server) = start_mock_sql_server(bodies.clone()).await; + let dir = tempfile::tempdir().unwrap(); + let command = ExportCommand::parse_from([ + "export", + "--addr", + client.addr(), + "--database", + "catalog-public", + "--output-dir", + dir.path().to_str().unwrap(), + "--target", + target, + "--auth-basic", + "user:password", + "--timeout", + "7s", + "--no-proxy", + ]); + let tool = command.build().await.unwrap(); + let result = tool.do_work().await; + drop(tool); + let queries = server.await.unwrap(); + assert!(result.unwrap_err().to_string().contains("middle")); + assert!( + dir.path() + .join("catalog/public/create_database.sql") + .is_file() + ); + assert!(!dir.path().join("catalog/public/create_tables.sql").exists()); + assert_eq!(queries.len(), 6); + assert!(queries[5].contains("SHOW CREATE TABLE \"catalog\".\"public\".\"middle\"")); + assert!(!queries.iter().any(|sql| sql.starts_with("COPY"))); + } + } + + #[tokio::test] + async fn sql_response_errors_do_not_expose_copy_credentials() { + let secret = "sentinel-connection-secret"; + let echoed = "sentinel-response-secret"; + let sql = format!( + "COPY DATABASE public TO 's3://bucket/' CONNECTION (SECRET_ACCESS_KEY='{secret}')" + ); + let bodies = vec![ + (200, format!("not json {echoed}")), + (200, json!({"output": [{"records": echoed}], "execution_time_ms": 0}).to_string()), + (400, json!({"code": 3001, "error": echoed}).to_string()), + (200, json!({"error": echoed}).to_string()), + (200, json!({"code": 3001, "output": [records("first", echoed)], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [{"error": echoed}]}).to_string()), + (200, json!({"output": [{"code": 3001, "records": records("first", echoed)["records"]}], "execution_time_ms": 0}).to_string()), + ]; + let count = bodies.len(); + let (client, server) = start_mock_sql_server(bodies).await; + for _ in 0..count { + let error = client + .sql_response(&sql, "public") + .await + .unwrap_err() + .to_string(); + assert!(error.contains("SQL"), "{error}"); + assert!(!error.contains(secret), "{error}"); + assert!(!error.contains(echoed), "{error}"); + assert!(!error.contains("COPY DATABASE"), "{error}"); + } + assert_eq!(server.await.unwrap().len(), count); + } + + #[tokio::test] + async fn rejects_http_and_statement_errors_without_retry() { + let valid = records("first", "CREATE TABLE first"); + let bodies = vec![ + (400, json!({"code": 3001, "error": "middle not found", "execution_time_ms": 0}).to_string()), + (200, json!({"code": 3001, "output": [valid.clone(), valid.clone(), valid.clone()], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [valid.clone(), {"error": "middle denied"}, valid.clone()], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [valid.clone(), {"error": "middle denied", "records": valid["records"]}, valid], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [records("first", "DDL"), {"code": 3001, "records": records("middle", "DDL")["records"]}, records("last", "DDL")], "execution_time_ms": 0}).to_string()), + (200, r#"{"execution_time_ms":0}"#.into()), + (200, "not json".into()), + ]; + let count = bodies.len(); + let (client, server) = start_mock_sql_server(bodies).await; + for _ in 0..count { + let mut ddl = String::new(); + let error = append_schema_ddl( + &client, + "catalog", + "schema", + ["first", "middle", "last"] + .into_iter() + .map(|name| ("TABLE", Some(name))), + &mut ddl, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("middle"), "{error}"); + assert!(ddl.is_empty()); + } + assert_eq!(server.await.unwrap().len(), count); + } +} diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index c5e9561d285..17008706066 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -295,8 +295,9 @@ pub trait SnapshotStorage: Send + Sync { schemas: &[String], chunk_id: u32, resume: bool, + packed: bool, ) -> Result<()> { - let _ = (schemas, chunk_id, resume); + let _ = (schemas, chunk_id, resume, packed); InvalidUriSnafu { uri: "snapshot", reason: "storage does not support preparing export chunks", @@ -730,6 +731,7 @@ impl SnapshotStorage for OpenDalStorage { schemas: &[String], chunk_id: u32, resume: bool, + packed: bool, ) -> Result<()> { let mut files = Vec::new(); for schema in schemas { @@ -751,10 +753,17 @@ impl SnapshotStorage for OpenDalStorage { continue; } let name = path.strip_prefix(&prefix).unwrap_or(""); - if !resume || entry.metadata().is_dir() || !valid_chunk_filename(name) { + if !resume + || entry.metadata().is_dir() + || !(if packed { + valid_packed_chunk_filename(name) + } else { + valid_chunk_filename(name) + }) + { return InvalidUriSnafu { uri: path, - reason: "expected an empty new chunk or direct Parquet files in an owned unfinished chunk", + reason: "expected an empty new chunk or recognized files in an owned unfinished chunk", }.fail(); } files.push(path.to_string()); @@ -786,6 +795,17 @@ impl SnapshotStorage for OpenDalStorage { } } +fn valid_packed_chunk_filename(name: &str) -> bool { + name == common_datasource::packed_snapshot::PACK_INDEX_FILE + || [("pack-", ".bin"), ("table-", ".parquet")] + .iter() + .any(|(prefix, suffix)| { + name.strip_prefix(prefix) + .and_then(|n| n.strip_suffix(suffix)) + .is_some_and(|id| !id.is_empty() && id.bytes().all(|b| b.is_ascii_digit())) + }) +} + fn valid_chunk_filename(name: &str) -> bool { name.strip_suffix(".parquet") .is_some_and(|stem| !stem.is_empty()) @@ -1074,7 +1094,7 @@ mod tests { let schemas = vec!["public".to_string(), "other".to_string()]; assert!( storage - .prepare_export_chunk(&schemas, 2, false) + .prepare_export_chunk(&schemas, 2, false, false) .await .is_err() ); @@ -1085,11 +1105,11 @@ mod tests { .unwrap() ); storage - .prepare_export_chunk(&schemas, 2, true) + .prepare_export_chunk(&schemas, 2, true, false) .await .unwrap(); storage - .prepare_export_chunk(&schemas, 2, true) + .prepare_export_chunk(&schemas, 2, true, false) .await .unwrap(); assert!( @@ -1129,7 +1149,7 @@ mod tests { } assert!( storage - .prepare_export_chunk(&["public".into(), "other".into()], 2, true) + .prepare_export_chunk(&["public".into(), "other".into()], 2, true, false) .await .is_err() ); @@ -1143,6 +1163,70 @@ mod tests { } } + #[tokio::test] + async fn packed_retry_preserves_foreign_files_and_completed_chunks() { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + let owned = ["pack-000000.bin", "table-42.parquet", "pack-index.json"]; + for name in owned { + storage + .write_text(&format!("data/public/2/{name}"), "partial") + .await + .unwrap(); + } + for path in [ + "data/public/1/pack-000000.bin", + "data/public/2/foreign.parquet", + ] { + storage.write_text(path, "keep").await.unwrap(); + } + let schemas = ["public".into()]; + assert!( + storage + .prepare_export_chunk(&schemas, 2, true, true) + .await + .is_err() + ); + for name in owned { + assert!( + storage + .file_exists(&format!("data/public/2/{name}")) + .await + .unwrap() + ); + } + assert_eq!( + storage + .read_text("data/public/2/foreign.parquet") + .await + .unwrap(), + "keep" + ); + storage + .object_store + .delete("data/public/2/foreign.parquet") + .await + .unwrap(); + storage + .prepare_export_chunk(&schemas, 2, true, true) + .await + .unwrap(); + assert!( + storage + .list_files_recursive("data/public/2/") + .await + .unwrap() + .is_empty() + ); + assert_eq!( + storage + .read_text("data/public/1/pack-000000.bin") + .await + .unwrap(), + "keep" + ); + } + #[tokio::test] async fn test_prepare_export_chunk_reports_delete_failure() { let dir = tempdir().unwrap(); @@ -1159,7 +1243,7 @@ mod tests { .await .unwrap(); let error = storage - .prepare_export_chunk(&["public".into()], 2, true) + .prepare_export_chunk(&["public".into()], 2, true, false) .await .unwrap_err(); let Error::StorageOperation { @@ -1191,7 +1275,7 @@ mod tests { .await .unwrap(); storage - .prepare_export_chunk(&[schema], 2, true) + .prepare_export_chunk(&[schema], 2, true, false) .await .unwrap(); assert!(!storage.file_exists(&owned).await.unwrap()); diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index b771ab48d05..1efe6d7d0a6 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.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::Duration; use base64::Engine; @@ -25,9 +26,7 @@ use servers::http::header::constants::GREPTIME_DB_HEADER_TIMEOUT; use servers::http::result::greptime_result_v1::GreptimedbV1Response; use snafu::ResultExt; -use crate::error::{ - BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result, SerdeJsonSnafu, -}; +use crate::error::{BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result}; #[derive(Debug, Clone)] pub struct DatabaseClient { @@ -37,6 +36,7 @@ pub struct DatabaseClient { timeout: Duration, proxy: Option, no_proxy: bool, + client: Arc>, } pub fn parse_proxy_opts( @@ -86,6 +86,7 @@ impl DatabaseClient { timeout, proxy, no_proxy, + client: Arc::new(tokio::sync::OnceCell::new()), } } @@ -97,48 +98,84 @@ impl DatabaseClient { self.sql(sql, DEFAULT_SCHEMA_NAME).await } - /// Requires the explicit packed-import protocol before any restore mutation. - pub async fn require_packed_import(&self) -> Result<()> { + pub(crate) fn catalog(&self) -> &str { + &self.catalog + } + + /// Requires packed export support before creating snapshot artifacts. + pub async fn require_packed_export(&self) -> Result<()> { + self.capabilities().await?.require("metric_packed_export") + } + + /// Reads target capabilities once; older targets advertise no extensions. + pub(crate) async fn capabilities(&self) -> Result { let url = format!("http://{}/v1/capabilities", self.addr); - let mut builder = reqwest::Client::builder().timeout(self.timeout); - if let Some(proxy) = self.proxy.clone() { - builder = builder.proxy(proxy); - } - if self.no_proxy { - builder = builder.no_proxy(); - } - let client = builder.build().context(BuildClientSnafu)?; - let mut request = client.get(&url); + let mut request = self.http_client().await?.get(&url).timeout(self.timeout); if let Some(auth) = &self.auth_header { request = request.header("Authorization", auth); } let response = request.send().await.with_context(|_| HttpQuerySqlSnafu { - reason: "packed import capability request failed", + reason: "capability request failed", })?; if response.status() == reqwest::StatusCode::NOT_FOUND { - return crate::error::InvalidArgumentsSnafu { - msg: "target does not support packed import", - } - .fail(); + return Ok(Capabilities(serde_json::json!({}))); } let response = response .error_for_status() .with_context(|_| HttpQuerySqlSnafu { - reason: "packed import capability request rejected", + reason: "capability request rejected", })?; let body = response.text().await.with_context(|_| HttpQuerySqlSnafu { reason: "cannot read capability response", })?; - let value: Value = serde_json::from_str(&body).context(SerdeJsonSnafu)?; - if !value.is_object() { + let value: Value = serde_json::from_str(&body).map_err(|_| { + crate::error::UnexpectedSnafu { + msg: "invalid capability response", + } + .build() + })?; + if !value.is_object() + || value.get("error").is_some() + || value.get("code").is_some_and(|v| v.as_u64() != Some(0)) + { return crate::error::UnexpectedSnafu { msg: "invalid capability response: expected JSON object", } .fail(); } - if value.get("metric_packed_import").and_then(Value::as_u64) != Some(1) { + for key in [ + "metric_batch_ddl", + "metric_packed_import", + "metric_packed_export", + ] { + if value.get(key).is_some_and(|v| v.as_u64().is_none()) { + return crate::error::UnexpectedSnafu { + msg: "invalid capability version", + } + .fail(); + } + } + Ok(Capabilities(value)) + } + + /// Submits a batch once. An ambiguous failure must not be replayed here. + pub(crate) async fn logical_tables(&self, sql: &str, schema: &str, count: usize) -> Result<()> { + let db = format!("{}-{}", self.catalog, schema); + if 8 + form_encoded_len(&db) + form_encoded_len(sql) > 4 * 1024 * 1024 { return crate::error::InvalidArgumentsSnafu { - msg: "target does not support packed import version 1", + msg: "logical-table batch exceeds 4 MiB form limit", + } + .fail(); + } + let response = self.ddl_response(sql, schema, "ddl/logical-tables").await?; + if response.output().len() != count + || response + .output() + .iter() + .any(|output| !matches!(output, GreptimeQueryOutput::AffectedRows(0))) + { + return crate::error::UnexpectedSnafu { + msg: "invalid logical-table batch output", } .fail(); } @@ -159,23 +196,39 @@ impl DatabaseClient { sql: &str, schema: &str, ) -> Result { - let url = format!("http://{}/v1/sql", self.addr); + self.ddl_response(sql, schema, "sql").await + } + + async fn ddl_response( + &self, + sql: &str, + schema: &str, + endpoint: &str, + ) -> Result { + let url = format!("http://{}/v1/{endpoint}", self.addr); let params = [ ("db", format!("{}-{}", self.catalog, schema)), ("sql", sql.to_string()), ]; - let mut builder = reqwest::Client::builder(); - if let Some(proxy) = self.proxy.clone() { - builder = builder.proxy(proxy); - } - if self.no_proxy { - builder = builder.no_proxy(); - } - let client = builder.build().context(BuildClientSnafu)?; - let mut request = client + let mut request = self + .http_client() + .await? .post(&url) - .form(¶ms) .header("Content-Type", "application/x-www-form-urlencoded"); + if endpoint == "ddl/logical-tables" { + let body = url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(¶ms) + .finish(); + // A streaming body cannot be cloned for reqwest's protocol retries. + request = request + .timeout(self.timeout) + .header(reqwest::header::CONTENT_LENGTH, body.len()) + .body(reqwest::Body::wrap_stream(futures::stream::once( + std::future::ready(Ok::<_, std::io::Error>(body)), + ))); + } else { + request = request.form(¶ms); + } if let Some(ref auth) = self.auth_header { request = request.header("Authorization", auth); } @@ -188,17 +241,105 @@ impl DatabaseClient { let response = request.send().await.with_context(|_| HttpQuerySqlSnafu { reason: format!("bad url: {}", url), })?; - let response = response - .error_for_status() - .with_context(|_| HttpQuerySqlSnafu { - reason: format!("query failed: {}", sql), - })?; - + let status = response.status(); let text = response.text().await.with_context(|_| HttpQuerySqlSnafu { reason: "cannot get response text".to_string(), })?; + let value: Value = serde_json::from_str(&text).map_err(|_| { + crate::error::UnexpectedSnafu { + msg: format!("invalid SQL response ({status})"), + } + .build() + })?; + if !status.is_success() + || value.get("error").is_some() + || value + .get("code") + .is_some_and(|code| code.as_u64() != Some(0)) + { + return crate::error::UnexpectedSnafu { + msg: format!( + "SQL request failed ({status}, code {:?})", + value.get("code").and_then(Value::as_u64) + ), + } + .fail(); + } + if let Some(outputs) = value.get("output").and_then(Value::as_array) { + for (index, output) in outputs.iter().enumerate() { + if output.get("error").is_some() + || output + .get("code") + .is_some_and(|code| code.as_u64() != Some(0)) + { + return crate::error::UnexpectedSnafu { + msg: format!( + "SQL statement {} failed (code {:?})", + index + 1, + output.get("code").and_then(Value::as_u64) + ), + } + .fail(); + } + } + } + serde_json::from_value(value).map_err(|_| { + crate::error::UnexpectedSnafu { + msg: format!("invalid SQL response ({status})"), + } + .build() + }) + } - serde_json::from_str::(&text).context(SerdeJsonSnafu) + async fn http_client(&self) -> Result<&reqwest::Client> { + self.client + .get_or_try_init(|| async { + let mut builder = reqwest::Client::builder().redirect( + reqwest::redirect::Policy::custom(|attempt| { + if attempt + .previous() + .first() + .is_some_and(|url| url.path() == "/v1/ddl/logical-tables") + { + attempt.stop() + } else { + reqwest::redirect::Policy::default().redirect(attempt) + } + }), + ); + if let Some(proxy) = self.proxy.clone() { + builder = builder.proxy(proxy); + } + if self.no_proxy { + builder = builder.no_proxy(); + } + builder.build().context(BuildClientSnafu) + }) + .await + } +} + +fn form_encoded_len(value: &str) -> usize { + url::form_urlencoded::byte_serialize(value.as_bytes()) + .map(str::len) + .sum() +} + +pub(crate) struct Capabilities(Value); + +impl Capabilities { + pub(crate) fn supports(&self, capability: &str) -> bool { + self.0.get(capability).and_then(Value::as_u64) == Some(1) + } + + pub(crate) fn require(&self, capability: &str) -> Result<()> { + if !self.supports(capability) { + return crate::error::InvalidArgumentsSnafu { + msg: format!("server does not support {capability} version 1"), + } + .fail(); + } + Ok(()) } } @@ -217,50 +358,234 @@ pub(crate) fn split_database(database: &str) -> Result<(String, Option)> } #[cfg(test)] -mod tests { +pub(crate) mod tests { #[tokio::test] async fn packed_capability_probe_checks_auth_and_protocol_version() { use tokio::io::{AsyncReadExt, AsyncWriteExt}; - for (status, body, supported) in [ - (200, r#"{"metric_packed_import":1,"future":2}"#, true), - (200, "{}", false), - (200, r#"{"metric_packed_import":2}"#, false), - (404, "{}", false), - (401, "{}", false), - (403, "{}", false), - (200, "[]", false), - (200, "invalid-json", false), - ] { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = vec![0; 4096]; - let n = socket.read(&mut request).await.unwrap(); - let request = String::from_utf8_lossy(&request[..n]).to_lowercase(); - assert!(request.starts_with("get /v1/capabilities ")); - assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza==")); - socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); - }); - let client = super::DatabaseClient::new( - address.to_string(), - "greptime".into(), - Some("user:password".into()), - std::time::Duration::from_secs(5), - None, - true, - ); - assert_eq!( - client.require_packed_import().await.is_ok(), - supported, - "{status}: {body}" - ); - server.await.unwrap(); + for export in [false, true] { + for (status, body, supported) in [ + (200, r#"{"metric_packed_import":1,"future":2}"#, true), + (200, "{}", false), + (200, r#"{"metric_packed_import":2}"#, false), + (404, "{}", false), + (401, "{}", false), + (403, "{}", false), + (200, "[]", false), + (200, "invalid-json", false), + ] { + let body = if export { + body.replace("metric_packed_import", "metric_packed_export") + } else { + body.to_string() + }; + let response_body = body.clone(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let body = response_body; + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let n = socket.read(&mut request).await.unwrap(); + let request = String::from_utf8_lossy(&request[..n]).to_lowercase(); + assert!(request.starts_with("get /v1/capabilities ")); + assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza==")); + socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + }); + let client = super::DatabaseClient::new( + address.to_string(), + "greptime".into(), + Some("user:password".into()), + std::time::Duration::from_secs(5), + None, + true, + ); + assert_eq!( + if export { + client.require_packed_export().await.is_ok() + } else { + client + .capabilities() + .await + .and_then(|c| c.require("metric_packed_import")) + .is_ok() + }, + supported, + "{status}: {body}" + ); + server.await.unwrap(); + } } } use super::*; + #[test] + fn form_size_counts_utf8_and_escaping() { + for (db, sql) in [ + ("greptime-public", "SELECT 1"), + ("测-试", "+&='秘密';\n"), + ("a_b.*", "a~b"), + ] { + let encoded = url::form_urlencoded::Serializer::new(String::new()) + .append_pair("db", db) + .append_pair("sql", sql) + .finish(); + assert_eq!( + 8 + form_encoded_len(db) + form_encoded_len(sql), + encoded.len() + ); + } + } + + #[tokio::test] + async fn batch_capabilities_distinguish_absence_from_errors() { + for (status, body, expected) in [ + (200, r#"{"metric_batch_ddl":1}"#, Some(true)), + (200, r#"{"metric_batch_ddl":2}"#, Some(false)), + (200, "{}", Some(false)), + (404, "", Some(false)), + (401, "{}", None), + (403, "{}", None), + (500, "{}", None), + (200, "[]", None), + (200, "bad", None), + (200, r#"{"metric_batch_ddl":"1"}"#, None), + (200, r#"{"error":"secret"}"#, None), + ] { + let (client, requests, server) = test_server(status, body).await; + let result = client.capabilities().await; + assert_eq!( + result.as_ref().ok().map(|c| c.supports("metric_batch_ddl")), + expected + ); + assert_eq!(requests.lock().unwrap().len(), 1); + server.abort(); + let _ = server.await; + } + } + + #[tokio::test] + async fn batch_redirects_are_not_replayed() { + for status in [303, 307, 308] { + let (client, requests, server) = test_server( + status, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + assert!( + client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .is_err() + ); + assert_eq!(requests.lock().unwrap().len(), 1); + assert!(requests.lock().unwrap()[0].starts_with("POST /v1/ddl/logical-tables ")); + // Ordinary SQL retains the shared client's existing redirect behavior. + assert!(client.sql("SELECT 1", "public").await.is_err()); + assert!(requests.lock().unwrap().len() > 2); + server.abort(); + let _ = server.await; + } + } + + #[tokio::test] + async fn batch_disconnect_and_timeout_are_errors() { + use tokio::io::AsyncReadExt; + for disconnect in [true, false] { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut client = DatabaseClient::new( + listener.local_addr().unwrap().to_string(), + "greptime".into(), + None, + Duration::from_millis(100), + None, + true, + ); + let (stop, stopped) = tokio::sync::oneshot::channel::<()>(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + if !disconnect { + let _ = stopped.await; + } + }); + assert!( + client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .is_err() + ); + let _ = stop.send(()); + server.await.unwrap(); + client.catalog = "测".repeat(500_000); + let error = client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .unwrap_err(); + assert!(format!("{error:?}").contains("4 MiB form limit")); + } + } + + pub(crate) async fn test_server( + status: u16, + body: &str, + ) -> ( + DatabaseClient, + Arc>>, + tokio::task::JoinHandle<()>, + ) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let client = DatabaseClient::new( + listener.local_addr().unwrap().to_string(), + "greptime".into(), + Some("user:password".into()), + Duration::from_secs(5), + None, + true, + ); + let requests = Arc::new(std::sync::Mutex::new(Vec::new())); + let captured = requests.clone(); + let body = body.to_string(); + let server = tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + loop { + let mut buf = [0; 4096]; + let n = socket.read(&mut buf).await.unwrap(); + if n == 0 { + break; + } + bytes.extend_from_slice(&buf[..n]); + if let Some(end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&bytes[..end]).to_lowercase(); + let length = headers + .lines() + .find_map(|l| l.strip_prefix("content-length: ")) + .map(|n| n.parse::().unwrap()) + .unwrap_or(0); + if bytes.len() >= end + 4 + length { + break; + } + } + } + captured + .lock() + .unwrap() + .push(String::from_utf8(bytes).unwrap()); + let location = if (300..400).contains(&status) { + "Location: /v1/sql\r\n" + } else { + "" + }; + socket.write_all(format!("HTTP/1.1 {status} Response\r\n{location}Content-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + } + }); + (client, requests, server) + } + #[test] fn test_split_database() { let result = split_database("catalog-schema").unwrap(); diff --git a/src/cli/src/metadata/snapshot.rs b/src/cli/src/metadata/snapshot.rs index 59b2599ad9c..9c010c2235b 100644 --- a/src/cli/src/metadata/snapshot.rs +++ b/src/cli/src/metadata/snapshot.rs @@ -12,11 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::path::{Component, Path, PathBuf}; + use async_trait::async_trait; use clap::{Parser, Subcommand}; use common_error::ext::BoxedError; use common_meta::snapshot::MetadataSnapshotManager; -use object_store::{ObjectStore, services}; +use object_store::ObjectStore; use crate::Tool; use crate::common::{ObjectStoreConfig, StoreConfig, new_fs_object_store}; @@ -270,25 +272,44 @@ fn build_object_store_and_resolve_file_path( fs_root: &str, file_path: &str, ) -> Result<(ObjectStore, String), BoxedError> { - let object_store = object_store.build().map_err(BoxedError::new)?; - let object_store = match object_store { - Some(object_store) => object_store, - None => new_fs_object_store(fs_root)?, - }; + if let Some(object_store) = object_store.build().map_err(BoxedError::new)? { + return Ok((object_store, file_path.to_string())); + } - let file_path = if object_store.info().scheme() == services::FS_SCHEME { - resolve_relative_path_with_current_dir(file_path).map_err(BoxedError::new)? + let absolute_path = + resolve_relative_path_with_current_dir(file_path).map_err(BoxedError::new)?; + let (root, key) = split_fs_snapshot_path(fs_root, &absolute_path); + Ok((new_fs_object_store(&root)?, key)) +} + +/// Separates a local absolute path into an operator root and a relative key. +fn split_fs_snapshot_path(fs_root: &str, absolute_path: &str) -> (String, String) { + let path = Path::new(absolute_path); + // OpenDAL keys cannot contain a Windows drive/UNC prefix or a root directory. + // Keep the filesystem anchor in the operator root, not in the object key. + let anchor: PathBuf = path + .components() + .take_while(|c| matches!(c, Component::Prefix(_) | Component::RootDir)) + .collect(); + let key = path + .components() + .filter(|c| !matches!(c, Component::Prefix(_) | Component::RootDir)) + .map(|c| c.as_os_str().to_string_lossy()) + .collect::>() + .join("/"); + // Preserve the existing custom-root semantics: cwd-resolved keys are rooted + // under --dir. The default root follows the input path's drive on Windows. + let root = if fs_root == "/" { + anchor.to_string_lossy() } else { - file_path.to_string() + fs_root.into() }; - - Ok((object_store, file_path)) + (root.into_owned(), key) } #[cfg(test)] mod tests { use std::env; - use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -311,44 +332,159 @@ mod tests { #[tokio::test] async fn test_cmd_resolve_file_path() { - common_telemetry::init_default_ut_logging(); - let cmd = RestoreCommand::parse_from([ - "", - "--file_name", - "metadata_snapshot.metadata.fb", - "--backend", - "memory-store", - "--store-addrs", - "memory://", - ]); - let tool = cmd.build().await.unwrap(); let current_dir = env::current_dir().unwrap(); - let file_path = current_dir.join("metadata_snapshot.metadata.fb"); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); + for file_path in [ + "metadata_snapshot.metadata.fb".to_string(), + current_dir + .join("metadata_snapshot.metadata.fb") + .display() + .to_string(), + ] { + let cmd = RestoreCommand::parse_from([ + "", + "--file_name", + &file_path, + "--backend", + "memory-store", + "--store-addrs", + "memory://", + ]); + let tool = cmd.build().await.unwrap(); + assert!(!Path::new(&tool.file_path).is_absolute()); + assert!( + Path::new(&tool.file_path) + .components() + .all(|c| matches!(c, Component::Normal(_))) + ); + let expected = current_dir + .components() + .filter_map(|c| match c { + Component::Normal(part) => Some(part.to_str().unwrap()), + _ => None, + }) + .chain(["metadata_snapshot.metadata.fb"]) + .collect::>() + .join("/"); + assert_eq!(tool.file_path, expected); + } + } + #[cfg(windows)] + #[test] + fn test_split_windows_snapshot_path() { + for (path, root, key) in [ + ( + r"C:\backup\snapshot.metadata.fb", + r"C:\", + "backup/snapshot.metadata.fb", + ), + ( + r"D:\备份 dir\snapshot.metadata.fb", + r"D:\", + "备份 dir/snapshot.metadata.fb", + ), + ( + r"\\server\share\backup\snapshot.metadata.fb", + r"\\server\share\", + "backup/snapshot.metadata.fb", + ), + ( + r"\\?\C:\backup\snapshot.metadata.fb", + r"\\?\C:\", + "backup/snapshot.metadata.fb", + ), + ] { + assert_eq!( + split_fs_snapshot_path("/", path), + (root.to_string(), key.to_string()) + ); + assert_eq!( + split_fs_snapshot_path(r"D:\root", path), + (r"D:\root".to_string(), key.to_string()) + ); + } + } + + #[tokio::test] + async fn test_fs_snapshot_paths() { + let temp_dir = tempfile::tempdir().unwrap(); + for name in ["snapshot.metadata.fb", "nested dir/快照.metadata.fb"] { + let path = temp_dir.path().join(name); + let (store, key) = build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + "/", + path.to_str().unwrap(), + ) + .unwrap(); + assert!(!path.exists()); + store.write(&key, "snapshot").await.unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), b"snapshot"); + assert_eq!(store.read(&key).await.unwrap().to_bytes(), &b"snapshot"[..]); + } + } + + #[tokio::test] + async fn test_fs_snapshot_custom_root() { + let temp_dir = tempfile::tempdir().unwrap(); + let (store, key) = build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + temp_dir.path().to_str().unwrap(), + "snapshot.metadata.fb", + ) + .unwrap(); + let current_dir = env::current_dir().unwrap(); + let expected = current_dir + .components() + .filter_map(|c| match c { + Component::Normal(part) => Some(part.to_str().unwrap()), + _ => None, + }) + .chain(["snapshot.metadata.fb"]) + .collect::>() + .join("/"); + assert_eq!(key, expected); + store.write(&key, "snapshot").await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join(expected)).unwrap(), + b"snapshot" + ); + } + + #[tokio::test] + async fn test_fs_snapshot_rejects_parent_traversal() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().join("../snapshot.metadata.fb"); + let (store, key) = build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + "/", + path.to_str().unwrap(), + ) + .unwrap(); + assert!(key.split('/').any(|part| part == "..")); + let err = store.write(&key, "snapshot").await.unwrap_err(); + assert_eq!(err.kind(), object_store::ErrorKind::NotFound); + assert!(err.to_string().contains("path escapes the configured root")); + } + + #[test] + fn test_remote_snapshot_key_unchanged() { let cmd = RestoreCommand::parse_from([ "", - "--file_name", - "metadata_snapshot.metadata.fb", "--backend", "memory-store", "--store-addrs", "memory://", + "--s3", + "--s3-bucket", + "snapshots", + "--s3-region", + "us-east-1", ]); - let tool = cmd.build().await.unwrap(); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); - - let cmd = RestoreCommand::parse_from([ - "", - "--file_name", - "metadata_snapshot.metadata.fb", - "--backend", - "memory-store", - "--store-addrs", - "memory://", - ]); - let tool = cmd.build().await.unwrap(); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); + let key = "nested/snapshot.metadata.fb"; + let (store, resolved) = + build_object_store_and_resolve_file_path(cmd.object_store, "/ignored", key).unwrap(); + assert_eq!(store.info().scheme(), object_store::services::S3_SCHEME); + assert_eq!(resolved, key); } async fn setup_backup_file(object_store: ObjectStore, file_path: &str) { diff --git a/src/client/Cargo.toml b/src/client/Cargo.toml index 5b479b1fbf2..36c3f60a9a4 100644 --- a/src/client/Cargo.toml +++ b/src/client/Cargo.toml @@ -45,7 +45,3 @@ substrait.workspace = true tokio.workspace = true tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true - -[dev-dependencies] -common-grpc-expr.workspace = true -tracing = "0.1" diff --git a/src/client/src/client.rs b/src/client/src/client.rs index fc40c98a723..dfc0e44fad5 100644 --- a/src/client/src/client.rs +++ b/src/client/src/client.rs @@ -87,6 +87,7 @@ struct Inner { impl Default for Inner { fn default() -> Self { + #[allow(deprecated)] Self::with_manager_and_peers(ChannelManager::new(), Vec::new(), ClientOptions::default()) } } @@ -105,6 +106,7 @@ struct Peers { } impl Inner { + #[deprecated(note = "legacy single-manager path shares its pool between lanes")] fn with_manager_and_peers( channel_manager: ChannelManager, peers: Vec, @@ -113,6 +115,7 @@ impl Inner { Self::with_managers_and_peers(channel_manager.clone(), channel_manager, peers, options) } + // The explicit dual-manager path keeps query and control pools independent. fn with_managers_and_peers( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, @@ -222,36 +225,61 @@ fn random_initial_delay(max_delay: Duration) -> Duration { } impl Client { + /// Creates a client whose query and control lanes intentionally share the default manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn new() -> Self { Default::default() } + /// Creates a client whose query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn with_urls(urls: A) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_urls_and_options(urls, ClientOptions::default()) } /// Creates a client with URLs and custom options. + /// + /// The query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_urls_and_options(urls: A, options: ClientOptions) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_manager_and_urls_and_options(ChannelManager::new(), urls, options) } + /// Creates a TLS client whose query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_tls_and_urls(urls: A, client_tls: ClientTlsOption) -> Result where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_tls_and_urls_and_options(urls, client_tls, ClientOptions::default()) } /// Creates a client with TLS URLs and custom options. + /// + /// The query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_tls_and_urls_and_options( urls: A, client_tls: ClientTlsOption, @@ -265,6 +293,7 @@ impl Client { let tls_config = load_client_tls_config(Some(client_tls)).context(error::CreateTlsChannelSnafu)?; let channel_manager = ChannelManager::with_config(channel_config, tls_config); + #[allow(deprecated)] Ok(Self::with_manager_and_urls_and_options( channel_manager, urls, @@ -272,15 +301,24 @@ impl Client { )) } + /// Creates a client with one manager shared intentionally by query and control lanes. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn with_manager_and_urls(channel_manager: ChannelManager, urls: A) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_manager_and_urls_and_options(channel_manager, urls, ClientOptions::default()) } - pub(crate) fn with_managers_and_urls( + /// Creates a client with query and control lanes backed by the supplied managers. + /// + /// The lanes are isolated only when the supplied managers are independently constructed; + /// this constructor does not enforce that they are distinct. + pub fn with_query_and_control_managers( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, urls: A, @@ -289,7 +327,7 @@ impl Client { U: AsRef, A: AsRef<[U]>, { - Self::with_managers_and_urls_and_options( + Self::with_query_and_control_managers_and_options( query_channel_manager, control_channel_manager, urls, @@ -297,7 +335,7 @@ impl Client { ) } - fn with_managers_and_urls_and_options( + fn with_query_and_control_managers_and_options( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, urls: A, @@ -323,6 +361,11 @@ impl Client { } /// Creates a client with a channel manager, URLs, and custom options. + /// + /// The query and control lanes intentionally share this manager and its pool. + #[deprecated( + note = "shares this manager and its pool between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_manager_and_urls_and_options( channel_manager: ChannelManager, urls: A, @@ -333,7 +376,8 @@ impl Client { A: AsRef<[U]>, { let channel_manager_for_query = channel_manager.clone(); - Self::with_managers_and_urls_and_options( + // Legacy constructors intentionally share the manager and therefore its pool. + Self::with_query_and_control_managers_and_options( channel_manager_for_query, channel_manager, urls, @@ -541,6 +585,23 @@ impl Client { Ok(()) } + /// Returns the number of cached channels in the query and control pools for tests. + #[cfg(feature = "testing")] + pub fn channel_pool_sizes(&self) -> (usize, usize) { + let pool_size = |channel_manager: &ChannelManager| { + let mut size = 0; + channel_manager.retain_channel(|_, _| { + size += 1; + true + }); + size + }; + ( + pool_size(&self.inner.query_channel_manager), + pool_size(&self.inner.control_channel_manager), + ) + } + /// Returns peer addresses grouped by active and inactive state for tests. #[cfg(feature = "testing")] pub fn peer_addresses_by_state(&self) -> (Vec, Vec) { @@ -559,6 +620,7 @@ impl Client { } #[cfg(test)] +#[allow(deprecated)] mod tests { use std::collections::HashSet; use std::sync::Arc; @@ -666,6 +728,36 @@ mod tests { ] } + #[tokio::test] + async fn test_explicit_dual_manager_constructor_uses_isolated_reused_channel_pools() { + let query_channel_manager = ChannelManager::new(); + let control_channel_manager = ChannelManager::new(); + let client = Client::with_query_and_control_managers( + query_channel_manager.clone(), + control_channel_manager.clone(), + ["127.0.0.1:3001"], + ); + + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_eq!(1, channel_pool_size(&query_channel_manager)); + assert_eq!(0, channel_pool_size(&control_channel_manager)); + + client.make_control_flight_client(false, false).unwrap(); + client.make_control_flight_client(false, false).unwrap(); + assert_eq!(1, channel_pool_size(&query_channel_manager)); + assert_eq!(1, channel_pool_size(&control_channel_manager)); + } + + fn channel_pool_size(manager: &ChannelManager) -> usize { + let mut size = 0; + manager.retain_channel(|_, _| { + size += 1; + true + }); + size + } + #[test] fn test_inner() { let inner = Inner::default(); diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs index 008ada00d23..4d723361e80 100644 --- a/src/client/src/client_manager.rs +++ b/src/client/src/client_manager.rs @@ -77,6 +77,7 @@ impl FlownodeManager for NodeClients { } impl NodeClients { + /// Creates node clients with independent query and control channel pools. pub fn new(config: ChannelConfig) -> Self { Self { query_channel_manager: ChannelManager::with_config(config.clone(), None), @@ -91,7 +92,7 @@ impl NodeClients { pub async fn get_client(&self, datanode: &Peer) -> Client { self.clients .get_with_by_ref(datanode, async move { - Client::with_managers_and_urls( + Client::with_query_and_control_managers( self.query_channel_manager.clone(), self.control_channel_manager.clone(), vec![datanode.addr.clone()], @@ -153,6 +154,7 @@ mod tests { assert_pool_has_one_address(&node_clients.control_channel_manager); let manager = ChannelManager::new(); + #[allow(deprecated)] let legacy = Client::with_manager_and_urls(manager.clone(), [PEER_ADDR]); legacy.make_flight_client(false, false).unwrap(); legacy.make_control_flight_client(false, false).unwrap(); diff --git a/src/client/src/database.rs b/src/client/src/database.rs index e97fff6db00..de3c722a86a 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, Mutex, RwLock}; use std::task::{Context, Poll}; use std::time::Duration; @@ -55,6 +55,7 @@ use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; use snafu::{IntoError, ResultExt}; +use tokio::sync::Notify; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; @@ -70,20 +71,47 @@ type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; const HINTS_METADATA_KEY: &str = "x-greptime-hints"; +/// Maximum time to wait for the optional trailing metrics message after +/// affected rows have already been delivered. +const FLIGHT_TRAILING_METRICS_TIMEOUT: Duration = Duration::from_secs(5); /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully -/// drained and [`Self::is_ready`] returns `true`. +/// drained and [`Self::is_ready`] returns `true`. Affected-row outputs may +/// briefly await a compatibility trailing metrics message. #[derive(Debug, Clone, Default)] pub struct OutputMetrics { inner: Arc, } -#[derive(Debug, Default)] +#[derive(Debug)] struct OutputMetricsInner { metrics: RwLock>, + completion_error: RwLock>, ready: AtomicBool, + ready_notify: Notify, + compatibility_task: Mutex>, +} + +impl Default for OutputMetricsInner { + fn default() -> Self { + Self { + metrics: RwLock::new(None), + completion_error: RwLock::new(None), + ready: AtomicBool::new(false), + ready_notify: Notify::new(), + compatibility_task: Mutex::new(None), + } + } +} + +impl Drop for OutputMetricsInner { + fn drop(&mut self) { + if let Some(handle) = self.compatibility_task.get_mut().unwrap().take() { + handle.abort(); + } + } } impl OutputMetrics { @@ -98,10 +126,45 @@ impl OutputMetrics { /// Marks the terminal metrics as final for this output. pub fn mark_ready(&self) { - let _ = self + if self .inner .ready - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire); + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.inner.ready_notify.notify_waiters(); + } + } + + /// Waits until terminal metrics are final. + pub async fn wait_ready(&self) { + loop { + let notified = self.inner.ready_notify.notified(); + if self.is_ready() { + return; + } + notified.await; + } + } + + /// Returns an error encountered while completing the output, if any. + pub fn completion_error(&self) -> Option { + self.inner.completion_error.read().unwrap().clone() + } + + fn set_completion_error(&self, error: impl Into) { + *self.inner.completion_error.write().unwrap() = Some(error.into()); + } + + fn set_compatibility_task(&self, handle: tokio::task::AbortHandle) { + let mut task = self.inner.compatibility_task.lock().unwrap(); + if !self.is_ready() { + *task = Some(handle); + } + } + + fn take_compatibility_task(&self) -> Option { + self.inner.compatibility_task.lock().unwrap().take() } /// Returns whether terminal metrics are final. @@ -148,7 +211,8 @@ impl OutputMetrics { /// /// The contained [`OutputMetrics`] lets callers read stream terminal metrics /// after consuming `output`. For non-stream outputs, metrics are ready -/// immediately. +/// immediately. Flight affected-row outputs without inline metrics require +/// [`OutputMetrics::wait_ready`] before compatibility trailing metrics can be read. #[derive(Debug)] pub struct OutputWithMetrics { pub output: Output, @@ -194,6 +258,55 @@ fn parse_terminal_metrics(metrics_json: &str) -> Result { }) } +fn spawn_affected_rows_trailing_metrics_task( + terminal_metrics: &OutputMetrics, + mut reader: FlightMessageReader, +) where + S: Stream> + Send + Unpin + 'static, +{ + let metrics_ref = Arc::downgrade(&terminal_metrics.inner); + let remote_addr = reader.remote_addr().to_string(); + let task = common_runtime::spawn_global(async move { + let result = + tokio::time::timeout(FLIGHT_TRAILING_METRICS_TIMEOUT, reader.read_next()).await; + let Some(inner) = metrics_ref.upgrade() else { + return; + }; + let metrics = OutputMetrics { inner }; + match result { + Ok(Ok(Some(FlightMessage::Metrics(s)))) => match parse_terminal_metrics(&s) { + Ok(metrics_json) => metrics.update(Some(metrics_json)), + Err(error) => { + metrics.set_completion_error(error.to_string()); + warn!( + "Failed to decode trailing Flight metrics from {}: {}", + remote_addr, error + ); + } + }, + Ok(Ok(None)) => {} + Ok(Ok(Some(other))) => { + let error = format!("Unexpected trailing Flight message: {other:?}"); + metrics.set_completion_error(error.clone()); + warn!("{} from {}", error, remote_addr); + } + Ok(Err(error)) => { + let error = flight_stream_error(&remote_addr, error); + metrics.set_completion_error(error.to_string()); + warn!("{}", error); + } + Err(_) => { + let error = "Timed out waiting for trailing Flight metrics"; + metrics.set_completion_error(error); + warn!("{} from {}", error, remote_addr); + } + } + metrics.mark_ready(); + metrics.take_compatibility_task(); + }); + terminal_metrics.set_compatibility_task(task.abort_handle()); +} + struct StreamWithMetrics { stream: common_recordbatch::SendableRecordBatchStream, metrics: OutputMetrics, @@ -233,6 +346,9 @@ impl Stream for StreamWithMetrics { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let polled = Pin::new(&mut self.stream).poll_next(cx); + if let Poll::Ready(Some(Err(error))) = &polled { + self.metrics.set_completion_error(error.to_string()); + } if let Poll::Ready(None) = &polled { self.sync_terminal_metrics(); self.metrics.mark_ready(); @@ -280,32 +396,12 @@ where FlightMessage::AffectedRows { rows, metrics } => { let terminal_metrics = OutputMetrics::new(); if let Some(metrics) = metrics { + // Inline metrics are authoritative. They complete the output + // without touching the rest of the Flight stream. terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); - } - let next_message = reader - .read_next() - .await - .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; - match next_message { - None => terminal_metrics.mark_ready(), - Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { - terminal_metrics.update(Some(parse_terminal_metrics(&s)?)); - terminal_metrics.mark_ready(); - } - Some(FlightMessage::Metrics(_)) => { - return IllegalFlightMessagesSnafu { - reason: "'AffectedRows' Flight metadata already carries Metrics and cannot be followed by another Metrics message", - } - .fail(); - } - Some(other) => { - return IllegalFlightMessagesSnafu { - reason: format!( - "'AffectedRows' Flight message can only be followed by a Metrics message, got {other:?}" - ), - } - .fail(); - } + terminal_metrics.mark_ready(); + } else { + spawn_affected_rows_trailing_metrics_task(&terminal_metrics, reader); } Ok(OutputWithMetrics { output: Output::new_with_affected_rows(rows), @@ -749,7 +845,9 @@ impl Database { /// Executes a SQL query and returns the output with terminal metrics. /// /// For stream outputs, callers must consume the stream before reading final - /// terminal metrics from [`OutputWithMetrics::metrics`]. + /// terminal metrics from [`OutputWithMetrics::metrics`]. For affected-row + /// outputs without inline metrics, call [`OutputMetrics::wait_ready`] when + /// compatibility trailing metrics are required. pub async fn sql_with_terminal_metrics( &self, sql: S, @@ -986,6 +1084,7 @@ mod tests { use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; use futures_util::StreamExt; + use tokio::sync::oneshot; use tonic::codegen::http::{HeaderMap, HeaderValue}; use tonic::metadata::MetadataMap; use tonic::{Code, Status}; @@ -1296,11 +1395,6 @@ mod tests { ); } - #[test] - fn test_parse_terminal_metrics_rejects_invalid_json() { - assert!(parse_terminal_metrics("{not-json}").is_err()); - } - #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { let output = output_from_flight_message_stream( @@ -1322,27 +1416,191 @@ mod tests { } #[tokio::test] - async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { + async fn test_affected_rows_inline_metrics_do_not_poll_trailer() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream( + let output = output_from_flight_message_stream( "test-peer".to_string(), futures_util::stream::iter(vec![ Ok(FlightMessage::AffectedRows { rows: 3, - metrics: Some(metrics_json.clone()), + metrics: Some(metrics_json), }), - Ok(FlightMessage::Metrics(metrics_json)), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(99))), ] as Vec>), ) .await - .unwrap_err(); + .unwrap(); - assert!( - err.to_string().contains("already carries Metrics"), - "unexpected error: {err:?}" + assert!(output.metrics.is_ready()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) ); } + #[tokio::test] + async fn test_affected_rows_without_inline_metrics_becomes_ready_after_trailer() { + // Hold the trailer until the not-ready state has been observed. + let (trailer_tx, trailer_rx) = oneshot::channel(); + let trailer = futures_util::stream::once(trailer_rx).map(|ready| { + ready.unwrap(); + Ok(FlightMessage::Metrics(terminal_metrics_json())) + }); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })]) + .chain(trailer), + ) + .await + .unwrap(); + + assert!(!output.metrics.is_ready()); + trailer_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(1), output.metrics.wait_ready()) + .await + .expect("terminal metrics must become ready once the trailer arrives"); + assert!(output.metrics.completion_error().is_none()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) + ); + } + + #[tokio::test] + async fn test_affected_rows_malformed_trailer_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("Invalid terminal metrics message")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_transport_trailer_error_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Err(Status::unavailable("trailer read failed").into()), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("trailer read failed")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_compatibility_reader_is_cancelled_after_second_poll_begins() { + struct DropProbe { + first: Option>, + polled: Option>, + dropped: Option>, + } + + impl Stream for DropProbe { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + if let Some(message) = self.first.take() { + return Poll::Ready(Some(message)); + } + if let Some(polled) = self.polled.take() { + let _ = polled.send(()); + } + Poll::Pending + } + } + + impl Drop for DropProbe { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } + } + + let (polled_tx, polled_rx) = oneshot::channel(); + let (dropped_tx, dropped_rx) = oneshot::channel(); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + DropProbe { + first: Some(Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })), + polled: Some(polled_tx), + dropped: Some(dropped_tx), + }, + ) + .await + .unwrap(); + polled_rx.await.unwrap(); + drop(output); + // Require cancellation before the five-second compatibility timeout. + tokio::time::timeout(Duration::from_secs(1), dropped_rx) + .await + .expect("the compatibility reader must be dropped once the output is dropped") + .unwrap(); + } + + #[tokio::test] + async fn test_schema_record_batch_yields_before_pending_message_stream() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "v", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let messages = futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ] as Vec>) + .chain(futures_util::stream::pending()); + let output = output_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + let OutputData::Stream(mut stream) = output.output.data else { + panic!("expected stream output"); + }; + + let batch = tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + assert!(!output.metrics.is_ready()); + } + #[tokio::test] async fn test_invalid_terminal_metrics_after_record_batch_yields_batch_then_error() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 1753fdd5246..3a21e6daa4e 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,40 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::pin::Pin; - use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use futures_util::stream::Peekable; use futures_util::{Stream, StreamExt}; use snafu::{OptionExt, ResultExt}; use crate::Result; use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum FlightMessageKind { - Schema, - RecordBatch, - AffectedRows, - Metrics, -} - -impl From<&FlightMessage> for FlightMessageKind { - fn from(message: &FlightMessage) -> Self { - match message { - FlightMessage::Schema(_) => Self::Schema, - FlightMessage::RecordBatch(_) => Self::RecordBatch, - FlightMessage::AffectedRows { .. } => Self::AffectedRows, - FlightMessage::Metrics(_) => Self::Metrics, - } - } -} - pub(crate) struct FlightMessageReader { /// Remote Flight peer associated with this response stream. remote_addr: String, - messages: Peekable, + messages: S, } impl FlightMessageReader @@ -55,7 +33,7 @@ where pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { Self { remote_addr: remote_addr.into(), - messages: messages.peekable(), + messages, } } @@ -72,21 +50,6 @@ where pub(crate) async fn read_next(&mut self) -> Result> { self.messages.next().await.transpose() } - - pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { - match Pin::new(&mut self.messages).peek().await { - Some(Ok(message)) => Ok(Some(message.into())), - None => Ok(None), - Some(Err(_)) => match self.read_next().await { - // `peek` only borrows the error; consume it to preserve the source error. - Err(error) => Err(error), - Ok(_) => IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek".to_string(), - } - .fail(), - }, - } - } } pub(crate) fn decode_flight_data( diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 4a4d243f36a..fb2e4d10ca1 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -48,7 +48,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); @@ -247,9 +247,7 @@ where "poll_flight_data_stream" )); - let mut stream_ended = false; - - while !stream_ended { + loop { let flight_message = match reader.read_next().await { Ok(Some(message)) => message, Ok(None) => break, @@ -262,59 +260,27 @@ where match flight_message { FlightMessage::RecordBatch(record_batch) => { - let result_to_yield = - RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - - // Metrics follow a batch so MergeScan can observe them before yielding it. - match reader.peek_next_message_kind().await { - Ok(Some(FlightMessageKind::Metrics)) => { - let metrics_message = match reader.read_next().await { - Ok(Some(FlightMessage::Metrics(metrics))) => metrics, - Ok(Some(_) | None) => { - yield IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek", - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(error) => { - yield Err(BoxedError::new(flight_stream_error( - &stream_addr, - error, - ))) - .context(ExternalSnafu); - break; - } - }; - let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); - metrics_ref.swap(metrics); - } - Ok(Some(FlightMessageKind::RecordBatch)) => {} - Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Ok(None) => stream_ended = true, - Err(error) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) - .context(ExternalSnafu); - break; - } - } - - yield Ok(result_to_yield); + // Deliver each batch immediately. In particular, do not + // wait for a possible following Metrics message; it is + // consumed on the next poll of this stream. + yield Ok(RecordBatch::from_df_record_batch( + schema_cloned.clone(), + record_batch, + )); } FlightMessage::Metrics(s) => { // Metrics may arrive before the next RecordBatch. - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); + match serde_json::from_str(&s) { + Ok(metrics) => { + metrics_ref.swap(Some(Arc::new(metrics))); + } + Err(error) => { + common_telemetry::warn!( + "Failed to decode region Flight metrics: {}", + error + ); + } + } continue; } _ => { @@ -418,6 +384,7 @@ pub fn check_response_header(header: &Option) -> Result<()> { } #[cfg(test)] +#[allow(deprecated)] mod test { use api::v1::Status as PbStatus; use api::v1::region::region_server::{Region, RegionServer}; @@ -718,7 +685,34 @@ mod test { } #[tokio::test] - async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + async fn test_record_batch_is_yielded_without_waiting_for_next_message() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let messages = stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ]) + .chain(stream::pending::>()); + let mut recordbatches = + recordbatches_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + + let batch = tokio::time::timeout(Duration::from_secs(1), recordbatches.next()) + .await + .expect("the first batch must not wait for lookahead") + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + } + + #[tokio::test] + async fn test_malformed_region_metrics_are_non_fatal() { let schema = test_schema(); let batch = RecordBatch::new( schema.clone(), @@ -729,23 +723,19 @@ mod test { "test-peer".to_string(), stream::iter(vec![ Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics(test_metrics_json())), ]), ) .await .unwrap(); - let batch = recordbatches.next().await.unwrap().unwrap(); - assert_eq!(batch.num_rows(), 1); - - let metrics = recordbatches.metrics().unwrap(); - assert_eq!(metrics.elapsed_compute, 7); + assert_eq!(recordbatches.next().await.unwrap().unwrap().num_rows(), 1); assert!(recordbatches.next().await.is_none()); } #[tokio::test] - async fn test_record_batch_stream_preserves_peeked_record_batch() { + async fn test_record_batch_stream_preserves_following_record_batch() { let schema = test_schema(); let first_batch = RecordBatch::new( schema.clone(), diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index fff268dea45..0ef2f35b371 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -22,6 +22,7 @@ required-features = ["dev-tools"] [features] default = [ "ai_functions", + "hdfs-object-store", "servers/pprof", "servers/mem-prof", "meta-srv/pg_kvbackend", @@ -32,9 +33,9 @@ enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterpr # Developer-only helper binaries and diagnostic datanode commands. # Kept out of `default` so normal/release builds don't compile them. dev-tools = [] +hdfs-object-store = ["mito2/hdfs-object-store", "object-store/hdfs-object-store"] mysql-object-store = ["object-store/mysql-object-store"] tokio-console = ["common-telemetry/tokio-console"] -vector_index = ["mito2/vector_index", "query/vector_index"] [lints] workspace = true @@ -43,10 +44,8 @@ workspace = true api.workspace = true async-trait.workspace = true auth.workspace = true -base64.workspace = true cache.workspace = true catalog.workspace = true -chrono.workspace = true clap.workspace = true cli.workspace = true client.workspace = true @@ -63,7 +62,6 @@ common-meta.workspace = true common-options.workspace = true common-procedure.workspace = true common-query.workspace = true -common-recordbatch.workspace = true common-runtime.workspace = true common-stat.workspace = true common-telemetry = { workspace = true, features = [ @@ -78,19 +76,16 @@ datafusion-physical-plan.workspace = true datafusion_object_store.workspace = true datanode.workspace = true datatypes.workspace = true -either = "1.15" etcd-client.workspace = true flow.workspace = true frontend = { workspace = true, default-features = false } futures.workspace = true hex.workspace = true human-panic = "2.0" -humantime.workspace = true lazy_static.workspace = true log-store.workspace = true meta-client.workspace = true meta-srv.workspace = true -metric-engine.workspace = true mito-codec.workspace = true mito2.workspace = true moka = { workspace = true, features = ["future"] } @@ -107,16 +102,13 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true -session.workspace = true smallvec.workspace = true snafu.workspace = true sqlparser.workspace = true standalone.workspace = true store-api.workspace = true -table.workspace = true tokio.workspace = true toml.workspace = true -tonic.workspace = true tracing-appender.workspace = true [target.'cfg(unix)'.dependencies] @@ -133,7 +125,9 @@ api.workspace = true client = { workspace = true, features = ["testing"] } common-test-util.workspace = true common-version.workspace = true +datanode = { workspace = true, features = ["testing"] } file-engine.workspace = true +metric-engine.workspace = true mito2.workspace = true serde.workspace = true similar-asserts.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 0f5239a6b51..271c3928e35 100644 --- a/src/cmd/src/bin/query_perf_fixture/case.rs +++ b/src/cmd/src/bin/query_perf_fixture/case.rs @@ -185,7 +185,7 @@ pub(super) fn default_max_concurrent_flushes() -> u64 { 256 } pub(super) fn default_worker_channel_capacity() -> u64 { - 65526 + 65_536 } pub(super) fn default_max_inflight_requests() -> u64 { 3000 diff --git a/src/cmd/src/cli.rs b/src/cmd/src/cli.rs index 09ec9646d8b..3739d4cc4fc 100644 --- a/src/cmd/src/cli.rs +++ b/src/cmd/src/cli.rs @@ -108,6 +108,7 @@ mod tests { use clap::Parser; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_grpc::channel_manager::ChannelManager; use common_telemetry::logging::LoggingOptions; use rand::Rng; @@ -174,7 +175,11 @@ mod tests { let mut instance = standalone.build(standalone_opts).await?; instance.start().await?; - let client = Client::with_urls([rpc_addr.as_str()]); + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [rpc_addr.as_str()], + ); let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); database .sql(r#"CREATE DATABASE "cli.export.create_table";"#) diff --git a/src/cmd/src/datanode/builder.rs b/src/cmd/src/datanode/builder.rs index 810dfeedc7b..6df7ac728e8 100644 --- a/src/cmd/src/datanode/builder.rs +++ b/src/cmd/src/datanode/builder.rs @@ -14,10 +14,9 @@ use std::sync::Arc; -use cache::build_datanode_cache_registry; +use cache::build_datanode_layered_cache_registry; use catalog::kvbackend::new_read_only_meta_kv_backend; use common_base::Plugins; -use common_meta::cache::LayeredCacheRegistryBuilder; use common_telemetry::info; use common_version::{short_version, verbose_version}; use datanode::datanode::DatanodeBuilder; @@ -111,11 +110,9 @@ impl InstanceBuilder { let backend = new_read_only_meta_kv_backend(client.clone()); let mut builder = DatanodeBuilder::new(dn_opts.clone(), plugins.clone(), backend.clone()); - let registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(backend)) - .build(), - ); + // The derived caches (table/partition caches) are put into the layer after the caches + // they are derived from, so they are invalidated only after them. + let registry = Arc::new(build_datanode_layered_cache_registry(backend)); builder .with_cache_registry(registry) .with_meta_client(client.clone()); diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index 14e289ea178..03696c96fc0 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -200,7 +200,6 @@ impl ObjbenchCommand { metadata: region_meta, source, cache_manager, - storage: None, max_sequence: None, sst_write_format: FormatType::PrimaryKey, index_options: Default::default(), @@ -209,8 +208,6 @@ impl ObjbenchCommand { fulltext_index_config, bloom_filter_index_config: MitoConfig::default().bloom_filter_index, preserve_row_sequence: false, - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; // Write SST diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index 6cae195f3e6..b000f7e1f02 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -467,8 +467,6 @@ impl StartCommand { // Some queries are expected to take long time. let mut channel_config = opts.datanode.client.channel_config(); channel_config.timeout = None; - // Source Flight streams and sink unary responses share pooled connections. - channel_config.http2_adaptive_window = Some(true); if opts.grpc.flight_compression.transport_compression() { channel_config.accept_compression = true; channel_config.send_compression = true; diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 4573b0f129f..76c95c26d32 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -44,14 +44,16 @@ use common_meta::procedure_executor::{LocalProcedureExecutor, ProcedureExecutorR use common_meta::region_keeper::MemoryRegionKeeper; use common_meta::region_registry::LeaderRegionRegistry; use common_meta::sequence::{Sequence, SequenceBuilder}; -use common_meta::wal_provider::{WalProviderRef, build_wal_provider}; +use common_meta::wal_provider::{WalProvider, WalProviderRef, build_wal_provider}; use common_options::plugin_options::StandaloneFlag; use common_procedure::ProcedureManagerRef; use common_query::prelude::set_default_prefix; -use common_telemetry::info; use common_telemetry::logging::{DEFAULT_LOGGING_DIR, TracingOptions}; +use common_telemetry::{info, warn}; use common_time::timezone::set_default_timezone; use common_version::{short_version, verbose_version}; +use common_wal::config::DatanodeWalConfig; +use common_wal::config::object_store::STANDALONE_GENERATION; use datanode::config::{DatanodeOptions, StorageConfig}; use datanode::datanode::{Datanode, DatanodeBuilder}; use datanode::region_server::RegionServer; @@ -83,6 +85,34 @@ use crate::{App, create_resource_limit_metrics, error, log_versions, maybe_activ pub const APP_NAME: &str = "greptime-standalone"; +/// Builds the WAL provider that allocates region WAL options in standalone mode. +/// +/// The object store WAL allocates the node prefix of `node_id`, the datanode id +/// of the standalone instance; an absent id is treated as 0, the id +/// `StandaloneOptions::datanode_options` assigns. +pub async fn build_standalone_wal_provider( + wal: &DatanodeWalConfig, + node_id: Option, + kv_backend: KvBackendRef, +) -> Result { + match wal { + DatanodeWalConfig::ObjectStore(config) => Ok(WalProvider::ObjectStore { + prefix: config.node_prefix(node_id.unwrap_or(0), STANDALONE_GENERATION), + }), + DatanodeWalConfig::RaftEngine(_) + | DatanodeWalConfig::Kafka(_) + | DatanodeWalConfig::Noop => { + let metasrv_wal_config = wal + .clone() + .try_into() + .context(error::InvalidWalProviderSnafu)?; + build_wal_provider(&metasrv_wal_config, kv_backend) + .await + .context(error::BuildWalProviderSnafu) + } + } +} + fn standalone_local_file_access( storage: &StorageConfig, ) -> common_datasource::error::Result { @@ -244,32 +274,53 @@ impl App for Instance { Ok(()) } + /// Stops every component in order. A failed step does not skip the later + /// ones; the first error is returned and the rest are logged. async fn stop(&mut self) -> Result<()> { - self.frontend - .shutdown() - .await - .context(error::ShutdownFrontendSnafu)?; + let mut first_error = None; + let mut record = |result: Result<()>| { + if let Err(err) = result { + if first_error.is_none() { + first_error = Some(err); + } else { + warn!(err; "Ignored a later shutdown error"); + } + } + }; - self.leader_services_controller - .stop( - self.procedure_manager.clone(), - self.datanode.region_server(), - ) - .await?; + record( + self.frontend + .shutdown() + .await + .context(error::ShutdownFrontendSnafu), + ); - self.datanode - .shutdown() - .await - .context(error::ShutdownDatanodeSnafu)?; + record( + self.leader_services_controller + .stop( + self.procedure_manager.clone(), + self.datanode.region_server(), + ) + .await, + ); - self.flownode - .shutdown() - .await - .context(error::ShutdownFlownodeSnafu)?; + record( + self.datanode + .shutdown() + .await + .context(error::ShutdownDatanodeSnafu), + ); + + record( + self.flownode + .shutdown() + .await + .context(error::ShutdownFlownodeSnafu), + ); info!("Datanode instance stopped."); - Ok(()) + first_error.map_or(Ok(()), Err) } } @@ -610,15 +661,8 @@ impl StartCommand { .step(10) .build(), ); - let kafka_options = opts - .wal - .clone() - .try_into() - .context(error::InvalidWalProviderSnafu)?; - let wal_provider = build_wal_provider(&kafka_options, kv_backend.clone()) - .await - .context(error::BuildWalProviderSnafu)?; - let wal_provider = Arc::new(wal_provider); + let wal_provider = + Arc::new(build_standalone_wal_provider(&opts.wal, node_id, kv_backend.clone()).await?); let table_metadata_allocator = Arc::new(TableMetadataAllocator::new( table_id_allocator.clone(), wal_provider.clone(), @@ -1035,16 +1079,136 @@ mod tests { use clap::{CommandFactory, Parser}; use common_base::readable_size::ReadableSize; use common_config::ENV_VAR_SEP; + use common_meta::ddl::allocator::wal_options::WalOptionsAllocator; + use common_meta::kv_backend::memory::MemoryKvBackend; use common_options::plugin_options::StandaloneFlag; use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; - use common_wal::config::DatanodeWalConfig; + use common_wal::config::object_store::ObjectStoreWalConfig; + use common_wal::options::{ObjectStoreWalOptions, WalOptions}; use frontend::frontend::FrontendOptions; use object_store::config::{FileConfig, GcsConfig}; use servers::grpc::GrpcOptions; + use servers::server::ServerHandlers; + use store_api::logstore::LogStore; use super::*; use crate::options::GlobalOptions; + struct FailingLeaderServices; + + #[async_trait] + impl StandaloneLeaderServicesController for FailingLeaderServices { + async fn start(&self, _context: LeaderServicesContext) -> Result<()> { + Ok(()) + } + + async fn stop( + &self, + _procedure_manager: ProcedureManagerRef, + _region_server: RegionServer, + ) -> Result<()> { + error::IllegalConfigSnafu { + msg: "leader services failed to stop", + } + .fail() + } + } + + struct FailingServer; + + #[async_trait] + impl servers::server::Server for FailingServer { + async fn shutdown(&self) -> servers::error::Result<()> { + servers::error::InternalSnafu { + err_msg: "server failed to stop", + } + .fail() + } + + async fn start(&mut self, _listening: SocketAddr) -> servers::error::Result<()> { + Ok(()) + } + + fn name(&self) -> &str { + "failing" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + #[tokio::test] + async fn test_stop_runs_every_step_and_returns_the_first_error() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("standalone-stop"); + let mut opts = StandaloneOptions { + wal: DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default()), + ..Default::default() + }; + opts.storage.data_home = data_home.path().to_str().unwrap().to_string(); + let creator = InstanceCreator::default() + .with_leader_services_controller(Box::new(FailingLeaderServices)); + let (mut instance, _) = StartCommand::build_with(opts, vec![], creator) + .await + .unwrap(); + let log_store = instance.datanode.object_store_log_store().unwrap(); + // The datanode fails to stop its servers after the leader services failed. + let mut services = ServerHandlers::default(); + services.insert((Box::new(FailingServer), "127.0.0.1:0".parse().unwrap())); + services.start_all().await.unwrap(); + instance.datanode.setup_services(services); + + let err = instance.stop().await.unwrap_err(); + + assert!( + matches!(&err, error::Error::IllegalConfig { msg, .. } if msg == "leader services failed to stop"), + "unexpected error: {err:?}" + ); + assert!(matches!( + log_store.append_batch(vec![]).await, + Err(log_store::error::Error::ObjectStoreWalStopped { .. }) + )); + } + + #[tokio::test] + async fn test_build_standalone_wal_provider() { + let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; + + let config = ObjectStoreWalConfig { + prefix: "cluster-a/wal".to_string(), + ..Default::default() + }; + // The regions persist the node prefix the datanode runs its store under. + let node_prefix = config.node_prefix(0, STANDALONE_GENERATION); + assert_eq!(node_prefix, "cluster-a/wal/datanodes/0/epochs/0"); + let provider = build_standalone_wal_provider( + &DatanodeWalConfig::ObjectStore(config), + Some(0), + kv_backend.clone(), + ) + .await + .unwrap(); + assert!(matches!( + &provider, + WalProvider::ObjectStore { prefix } if *prefix == node_prefix + )); + let regions = vec![0, 1]; + let wal_options = provider.allocate(®ions, false).await.unwrap(); + for region in regions { + assert_eq!( + wal_options[®ion], + WalOptions::ObjectStore(ObjectStoreWalOptions::new(node_prefix.clone())) + ); + } + + let provider = + build_standalone_wal_provider(&DatanodeWalConfig::default(), Some(0), kv_backend) + .await + .unwrap(); + assert!(matches!(provider, WalProvider::RaftEngine)); + } + #[test] fn test_standalone_local_file_access_config() { let data_home = create_temp_dir("standalone_copy_root"); @@ -1139,7 +1303,6 @@ mod tests { fn test_toml() { let opts = StandaloneOptions::default(); let toml_string = toml::to_string(&opts).unwrap(); - assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); let parsed: StandaloneOptions = toml::from_str(&toml_string).unwrap(); assert_eq!(parsed.otlp, opts.otlp); } diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 67e2e573f38..00cdafd9b8a 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -30,6 +30,7 @@ use datanode::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use file_engine::config::EngineConfig as FileEngineConfig; use flow::FlownodeOptions; use frontend::frontend::FrontendOptions; +use frontend::service_config::PendingRowsBatcherOptions; use meta_client::MetaClientOptions; use meta_srv::metasrv::MetasrvOptions; use meta_srv::selector::SelectorType; @@ -124,6 +125,44 @@ fn test_load_runtime_options_without_max_blocking_threads() { ); } +#[test] +#[cfg(feature = "hdfs-object-store")] +fn test_load_datanode_hdfs_config() { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + config.path(), + r#" + [storage] + type = "Hdfs" + name = "local-hdfs" + root = "/greptimedb" + name_node = "hdfs://127.0.0.1:9000" + enable_read_cache = false + options = { "dfs.client.block.write.replace-datanode-on-failure.enable" = "true" } + "#, + ) + .unwrap(); + + let options = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let object_store::config::ObjectStoreConfig::Hdfs(hdfs) = options.component.storage.store + else { + unreachable!() + }; + + assert_eq!("local-hdfs", hdfs.name); + assert_eq!("/greptimedb", hdfs.connection.root); + assert_eq!("hdfs://127.0.0.1:9000", hdfs.connection.name_node); + assert_eq!( + Some(&"true".to_string()), + hdfs.connection + .options + .get("dfs.client.block.write.replace-datanode-on-failure.enable") + ); + assert!(!hdfs.cache.enable_read_cache); +} + #[allow(deprecated)] #[test] fn test_load_datanode_example_config() { @@ -158,6 +197,7 @@ fn test_load_datanode_example_config() { }, region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { + experimental_series_index_max_size: ReadableSize::gb(5), auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), @@ -198,14 +238,12 @@ fn test_load_frontend_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); - assert!( - !options - .component - .otlp - .experimental_enable_exponential_histogram - ); let expected = GreptimeOptions:: { component: FrontendOptions { + pending_rows_batcher: PendingRowsBatcherOptions { + logical_table: Some(Default::default()), + ..Default::default() + }, default_timezone: Some("UTC".to_string()), default_column_prefix: Some("greptime".to_string()), auto_create_table: true, @@ -381,14 +419,12 @@ fn test_load_standalone_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); - assert!( - !options - .component - .otlp - .experimental_enable_exponential_histogram - ); let expected = GreptimeOptions:: { component: StandaloneOptions { + pending_rows_batcher: PendingRowsBatcherOptions { + logical_table: Some(Default::default()), + ..Default::default() + }, default_timezone: Some("UTC".to_string()), default_column_prefix: Some("greptime".to_string()), auto_create_table: true, @@ -400,6 +436,7 @@ fn test_load_standalone_example_config() { }), region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { + experimental_series_index_max_size: ReadableSize::gb(5), auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), @@ -427,6 +464,7 @@ fn test_load_standalone_example_config() { cors_allowed_origins: vec!["https://example.com".to_string()], ..Default::default() }, + grpc: GrpcOptions::default(), query: QueryOptions { memory_pool_size: MemoryLimit::Percentage(50), ..Default::default() @@ -438,41 +476,313 @@ fn test_load_standalone_example_config() { similar_asserts::assert_eq!(options, expected); } -#[test] -fn test_load_otlp_exponential_histogram_option() { - let config = tempfile::NamedTempFile::new().unwrap(); - std::fs::write( - config.path(), - "[otlp]\nexperimental_enable_exponential_histogram = true\n", - ) +/// Renders an example config with its `#+` options enabled, as `config.md` +/// lists them. With `drop_none_default`, options annotated with +/// `@toml2docs:none-default` are removed, because `config.md` does not present +/// their example value as the default. +fn render_example(example: &str, drop_none_default: bool) -> String { + let mut rendered = String::new(); + let mut none_default = false; + for line in example.lines() { + let line = line.trim_start(); + if line.starts_with("##") { + none_default |= line.contains("@toml2docs:none-default"); + continue; + } + let line = line.strip_prefix("#+").map_or(line, str::trim_start); + let is_option = !line.starts_with('[') && line.contains('='); + if !(is_option && drop_none_default && none_default) { + rendered.push_str(line); + rendered.push('\n'); + } + none_default = false; + } + rendered +} + +/// Flattens leaf values into dotted paths. Arrays of tables such as +/// `region_engine` are merged, since each element is keyed by its engine name. +fn leaf_values(value: &toml::Value, prefix: &str, leaves: &mut Vec<(String, toml::Value)>) { + match value { + toml::Value::Table(table) => { + for (key, value) in table { + let path = if prefix.is_empty() { + key.clone() + } else { + format!("{prefix}.{key}") + }; + leaf_values(value, &path, leaves); + } + } + toml::Value::Array(array) + if array.iter().all(toml::Value::is_table) && !array.is_empty() => + { + for element in array { + leaf_values(element, prefix, leaves); + } + } + _ => leaves.push((prefix.to_string(), value.clone())), + } +} + +/// Removes the option at a dotted path produced by [`leaf_values`]. +fn remove_leaf(value: &mut toml::Value, path: &str) -> bool { + match value { + toml::Value::Array(array) => array.iter_mut().any(|element| remove_leaf(element, path)), + toml::Value::Table(table) => match path.split_once('.') { + None => table.remove(path).is_some(), + Some((head, rest)) => table + .get_mut(head) + .is_some_and(|value| remove_leaf(value, rest)), + }, + _ => false, + } +} + +fn load_as_leaves( + config: Option<&str>, +) -> common_config::error::Result> +where + T: Configurable + serde::Serialize, +{ + let file = config.map(|config| { + let mut file = create_named_temp_file(); + file.write_all(config.as_bytes()).unwrap(); + file + }); + let options = GreptimeOptions::::load_layered_options( + file.as_ref().map(|file| file.path().to_str().unwrap()), + "", + )?; + let mut leaves = Vec::new(); + leaf_values(&toml::Value::try_from(&options).unwrap(), "", &mut leaves); + Ok(leaves.into_iter().collect()) +} + +/// Returns the default of an option inside the example's other settings, so +/// that options in sections absent by default still have one to compare with. +/// Removes the closest enclosing table that still loads without the option. +fn default_within_example(full: &toml::Value, path: &str) -> Option +where + T: Configurable + serde::Serialize, +{ + let mut removed = path; + loop { + let mut without = full.clone(); + assert!(remove_leaf(&mut without, removed)); + if let Ok(defaults) = load_as_leaves::(Some(&toml::to_string(&without).unwrap())) { + return defaults.get(path).cloned(); + } + removed = removed + .rsplit_once('.') + .unwrap_or_else(|| panic!("the example does not load without `{path}`")) + .0; + } +} + +/// Checks that every option in an example config exists, and that every value +/// `config.md` presents as a default is the actual default. +/// +/// `variants` lists tagged sections, e.g. `("wal", "provider", ["kafka"])`, +/// whose options only deserialize under a specific tag value. +fn check_example_defaults(name: &str, variants: &[(&str, &str, &[&str])]) -> Vec +where + T: Configurable + serde::Serialize, +{ + let example = std::fs::read_to_string(common_test_util::find_workspace_path(&format!( + "config/{name}" + ))) .unwrap(); + let full: toml::Value = toml::from_str(&render_example(&example, false)).unwrap(); + let documented: toml::Value = toml::from_str(&render_example(&example, true)).unwrap(); - let frontend = - GreptimeOptions::::load_layered_options(config.path().to_str(), "") - .unwrap(); - assert!( - frontend - .component - .otlp - .experimental_enable_exponential_histogram - ); + // (options loaded from the example, options loaded without config) per variant + let mut loaded = vec![( + load_as_leaves::(Some(&toml::to_string(&full).unwrap())).unwrap(), + load_as_leaves::(None).unwrap(), + )]; + for (section, tag, values) in variants { + for value in *values { + let mut variant = full.clone(); + variant[section][tag] = toml::Value::String(value.to_string()); + loaded.push(( + load_as_leaves::(Some(&toml::to_string(&variant).unwrap())).unwrap(), + load_as_leaves::(Some(&format!("[{section}]\n{tag} = \"{value}\"\n"))).unwrap(), + )); + } + } - let standalone = - GreptimeOptions::::load_layered_options(config.path().to_str(), "") - .unwrap(); - assert!( - standalone - .component - .otlp - .experimental_enable_exponential_histogram - ); - assert!( - standalone - .component - .frontend_options() - .otlp - .experimental_enable_exponential_histogram - ); + // Deserialize-only fields, invisible in the serialized options. + const SKIP_SERIALIZING: &[&str] = &[ + "storage.access_key_id", + "storage.secret_access_key", + "storage.access_key_secret", + "storage.account_name", + "storage.account_key", + "storage.credential_path", + "storage.credential", + "region_engine.mito.inverted_index.intermediate_path", + ]; + + let mut errors = Vec::new(); + let mut full_leaves = Vec::new(); + leaf_values(&full, "", &mut full_leaves); + for (path, _) in full_leaves { + if SKIP_SERIALIZING.contains(&path.as_str()) + || (path == "tracing.tokio_console_addr" && !cfg!(feature = "tokio-console")) + { + continue; + } + if !loaded + .iter() + .any(|(example, _)| example.contains_key(&path)) + { + errors.push(format!("{name}: `{path}` is not a known option")); + } + } + + // Options left unset by default and resolved where they are used. + let resolved_at_runtime: [(&str, toml::Value); 7] = [ + ( + "logging.dir", + format!("{DEFAULT_DATA_HOME}/{DEFAULT_LOGGING_DIR}").into(), + ), + ("logging.otlp_endpoint", DEFAULT_OTLP_HTTP_ENDPOINT.into()), + ("logging.otlp_export_protocol", "http".into()), + ("logging.tracing_sample_ratio.default_ratio", 1.0.into()), + ("wal.sync_period", "5s".into()), + ( + "storage.copy_root", + format!("{DEFAULT_DATA_HOME}/copy").into(), + ), + ("flow.batching_mode.frontend_tls.enabled", false.into()), + ]; + + let mut documented_leaves = Vec::new(); + leaf_values(&documented, "", &mut documented_leaves); + for (path, documented_value) in documented_leaves { + if SKIP_SERIALIZING.contains(&path.as_str()) { + continue; + } + if let Some((_, resolved)) = resolved_at_runtime.iter().find(|(p, _)| *p == path) { + if &documented_value != resolved { + errors.push(format!( + "{name}: `{path}` is documented as default `{documented_value}`, but resolves to `{resolved}`" + )); + } + continue; + } + let mut pairs = loaded + .iter() + .filter_map(|(example, default)| Some((example.get(&path)?, default.get(&path)?))) + .map(|(example, default)| (example.clone(), default.clone())) + .collect::>(); + if pairs.is_empty() { + // The option sits in a section that is absent without config, or + // is itself optional. + match default_within_example::(&full, &path) { + Some(default) => pairs.push((loaded[0].0[&path].clone(), default)), + None => { + errors.push(format!( + "{name}: `{path}` has no default to compare with; annotate it with \ + `@toml2docs:none-default` or list its runtime default in this test" + )); + continue; + } + } + } + if pairs.iter().all(|(example, default)| example != default) { + let (example, default) = &pairs[0]; + errors.push(format!( + "{name}: `{path}` is documented as default `{example}`, but the default is `{default}`; \ + fix the value or annotate it with `@toml2docs:none-default`" + )); + } + } + errors +} + +// The examples document HDFS storage options. +#[cfg(feature = "hdfs-object-store")] +#[test] +fn test_example_configs_document_actual_defaults() { + const WAL: (&str, &str, &[&str]) = ("wal", "provider", &["kafka", "experimental_object_store"]); + const STORAGE: (&str, &str, &[&str]) = + ("storage", "type", &["S3", "Oss", "Azblob", "Gcs", "Hdfs"]); + + let mut errors = + check_example_defaults::("standalone.example.toml", &[WAL, STORAGE]); + errors.extend(check_example_defaults::( + "frontend.example.toml", + &[], + )); + errors.extend(check_example_defaults::( + "datanode.example.toml", + &[WAL, STORAGE], + )); + errors.extend(check_example_defaults::( + "metasrv.example.toml", + &[("wal", "provider", &["kafka"])], + )); + errors.extend(check_example_defaults::( + "flownode.example.toml", + &[], + )); + assert!(errors.is_empty(), "{}", errors.join("\n")); +} + +#[test] +fn test_load_trace_aux_cache_size() { + let default = FrontendOptions::default().otlp.trace_aux_cache_size; + for (setting, size) in [ + ("", default), + ("trace_aux_cache_size = 0", ReadableSize(0)), + ("trace_aux_cache_size = \"1MiB\"", ReadableSize::mb(1)), + ] { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(config.path(), format!("[otlp]\n{setting}\n")).unwrap(); + + let frontend = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let standalone = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + for options in [frontend.component, standalone.component.frontend_options()] { + assert_eq!(options.otlp.trace_aux_cache_size, size); + } + } +} + +#[test] +fn test_load_removed_histogram_options() { + for enabled in [false, true] { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + config.path(), + format!( + "[otlp]\nexperimental_enable_exponential_histogram = {enabled}\n\ + [prom_store]\nenable = true\nwith_metric_engine = true\n\ + experimental_enable_prometheus_native_histogram = {enabled}\n" + ), + ) + .unwrap(); + + let frontend = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let standalone = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let defaults = FrontendOptions::default(); + for options in [frontend.component, standalone.component.frontend_options()] { + assert_eq!(options.otlp, defaults.otlp); + assert_eq!(options.prom_store, defaults.prom_store); + let serialized = toml::to_string(&options).unwrap(); + assert!(!serialized.contains("experimental_enable_exponential_histogram")); + assert!(!serialized.contains("experimental_enable_prometheus_native_histogram")); + } + } } #[test] diff --git a/src/common/base/Cargo.toml b/src/common/base/Cargo.toml index 760241a3e6a..702171e9d67 100644 --- a/src/common/base/Cargo.toml +++ b/src/common/base/Cargo.toml @@ -4,9 +4,6 @@ version.workspace = true edition.workspace = true license.workspace = true -[features] -testing = [] - [lints] workspace = true diff --git a/src/common/config/Cargo.toml b/src/common/config/Cargo.toml index 27b238add73..ad356716867 100644 --- a/src/common/config/Cargo.toml +++ b/src/common/config/Cargo.toml @@ -29,4 +29,3 @@ datanode.workspace = true meta-client.workspace = true serde.workspace = true temp-env = "0.3" -tempfile.workspace = true diff --git a/src/common/config/src/file_watcher.rs b/src/common/config/src/file_watcher.rs index d0766920191..da2c84a8f60 100644 --- a/src/common/config/src/file_watcher.rs +++ b/src/common/config/src/file_watcher.rs @@ -21,9 +21,16 @@ //! file itself. This ensures that file deletions and recreations are properly //! tracked, which is common with editors that use atomic saves or when //! configuration files are replaced. +//! +//! If the path goes through symlinks (the file itself or any ancestor +//! directory), the directory holding each symlink is watched as well, so that +//! swapping a symlink anywhere in the chain (e.g. `ln -sfn secrets.2 secrets`, +//! as done by secret managers like sops-nix or Kubernetes volume mounts) is +//! detected. The chain is re-resolved after every relevant event and the set +//! of watched directories is updated accordingly. -use std::collections::HashSet; -use std::path::{Path, PathBuf}; +use std::collections::{HashSet, VecDeque}; +use std::path::{Component, Path, PathBuf}; use std::sync::mpsc::channel; use common_telemetry::{error, info, warn}; @@ -54,10 +61,12 @@ impl FileWatcherConfig { /// /// The watcher monitors the parent directory of each file to handle file /// deletion and recreation properly. Events are filtered to only trigger -/// callbacks for the specific files being watched. +/// callbacks for the specific files being watched (or the symlinks leading +/// to them). pub struct FileWatcherBuilder { config: FileWatcherConfig, - /// Canonicalized paths of files to watch. + /// Paths of files to watch, as configured. Symlinks in them are resolved + /// when spawning and again after every relevant event. file_paths: Vec, } @@ -79,7 +88,7 @@ impl FileWatcherBuilder { /// Add a file path to watch. /// /// Returns an error if the path is a directory. - /// The path is canonicalized for reliable comparison with events. + /// Symlinks in the path are followed, see [`FileWatcherBuilder::spawn`]. pub fn watch_path>(mut self, path: P) -> Result { let path = path.as_ref(); snafu::ensure!( @@ -110,7 +119,10 @@ impl FileWatcherBuilder { /// /// The callback is invoked when relevant file events are detected for /// the watched files. The watcher monitors the parent directories to - /// handle file deletion and recreation properly. + /// handle file deletion and recreation properly. When a path goes through + /// symlinks, the directory holding each symlink is watched too, and the + /// watches are re-armed after every relevant event, so the callback also + /// fires when a symlink anywhere in the chain is replaced. /// /// The spawned watcher thread runs for the lifetime of the process. pub fn spawn(self, callback: F) -> Result<()> @@ -121,42 +133,61 @@ impl FileWatcherBuilder { let mut watcher = notify::recommended_watcher(tx).context(FileWatchSnafu { path: "" })?; - // Collect unique parent directories to watch - let mut watched_dirs: HashSet = HashSet::new(); - for file_path in &self.file_paths { - if let Some(parent) = file_path.parent() - && watched_dirs.insert(parent.to_path_buf()) - { - watcher - .watch(parent, RecursiveMode::NonRecursive) - .context(FileWatchSnafu { - path: parent.display().to_string(), - })?; + let mut targets = resolve_watch_targets(&self.file_paths); + let mut skipped = Vec::new(); + for dir in &targets.dirs { + match watcher.watch(dir, RecursiveMode::NonRecursive) { + Ok(()) => warn_if_root(dir), + // Only the directories holding the files are required. A + // directory holding a symlink may be unreadable (e.g. mode + // 0711); it is skipped and retried by `rearm_watches`. + Err(err) if !targets.file_dirs.contains(dir) => { + warn!("Failed to watch {:?}, skipping: {}", dir, err); + skipped.push(dir.clone()); + } + Err(err) => { + return Err(err).context(FileWatchSnafu { + path: dir.display().to_string(), + }); + } } } + for dir in skipped { + targets.dirs.remove(&dir); + } let config = self.config; + let file_paths = self.file_paths; info!( - "Spawning file watcher for paths: {:?} (watching parent directories)", - self.file_paths + "Spawning file watcher for paths: {:?} (watching directories: {:?})", + file_paths .iter() .map(|p| p.display().to_string()) - .collect::>() + .collect::>(), + targets.dirs ); std::thread::spawn(move || { // Keep watcher alive in the thread - let _watcher = watcher; + let mut watcher = watcher; while let Ok(res) = rx.recv() { match res { Ok(event) => { - if !is_relevant_event(&event.kind, &config) { + if !is_relevant_event(&event.kind, &config) + || !event + .paths + .iter() + .any(|p| targets.entries.contains(&entry_key(p))) + { continue; } info!(?event.kind, ?event.paths, "Detected folder change"); + // A symlink in the chain may have been replaced, so + // re-resolve and re-arm the watches before reloading. + targets = rearm_watches(&mut watcher, &file_paths, targets); callback(); } Err(err) => { @@ -178,6 +209,166 @@ impl Default for FileWatcherBuilder { } } +/// Maximum number of symlinks followed while resolving a single path. This +/// mirrors Linux's `MAXSYMLINKS` and guards against symlink loops. +const MAX_SYMLINK_HOPS: usize = 40; + +/// Directories to watch and the directory entries whose events are relevant. +#[derive(Debug, Default)] +struct WatchTargets { + /// Directories holding a symlink of a chain, or the final file. + dirs: HashSet, + /// The subset of `dirs` holding the final files. + file_dirs: HashSet, + /// The symlinks of every chain and the final files, as they appear + /// inside `dirs`, normalized with [`entry_key`]. + entries: HashSet, +} + +/// Key used to compare an event path with [`WatchTargets::entries`]. +/// +/// The entries are built from the configured spelling, which may differ in +/// case from what is on disk (and reported in events) on the usually +/// case-insensitive filesystems of macOS and Windows, so compare lowercased +/// there. At worst this lets a few extra events through on a case-sensitive +/// volume. +fn entry_key(path: &Path) -> PathBuf { + #[cfg(any(target_os = "macos", windows))] + { + PathBuf::from(path.to_string_lossy().to_lowercase()) + } + #[cfg(not(any(target_os = "macos", windows)))] + { + path.to_path_buf() + } +} + +/// Resolves every path component by component, following symlinks, and +/// collects what has to be watched to notice any change of what the path +/// points to. +/// +/// For each symlink met on the way (the file itself or any ancestor +/// directory), the directory holding it is watched, so replacing the symlink +/// is noticed. The directory holding the final file is watched too, so +/// in-place edits are noticed. The collected directories are free of +/// symlinks; the collected entries are the symlinks themselves and the final +/// files, located inside those directories. +/// +/// If a component is missing (e.g. a symlink being swapped non-atomically), +/// resolution of that path stops and its containing directory is watched, so +/// the watch is re-armed once the component shows up again. +fn resolve_watch_targets(paths: &[PathBuf]) -> WatchTargets { + let mut targets = WatchTargets::default(); + for path in paths { + resolve_one(path, &mut targets); + } + targets +} + +fn resolve_one(path: &Path, targets: &mut WatchTargets) { + let Ok(path) = std::path::absolute(path) else { + return; + }; + + let mut pending: VecDeque = path + .components() + .map(|c| PathBuf::from(c.as_os_str())) + .collect(); + // Always free of symlinks, since each component is resolved before + // being appended. + let mut resolved = PathBuf::new(); + let mut hops = 0; + + while let Some(component) = pending.pop_front() { + match component.components().next() { + Some(Component::Prefix(_) | Component::RootDir) => resolved.push(&component), + Some(Component::CurDir) | None => {} + Some(Component::ParentDir) => { + resolved.pop(); + } + Some(Component::Normal(name)) => { + let candidate = resolved.join(name); + let is_last = pending.is_empty(); + match std::fs::symlink_metadata(&candidate) { + Ok(meta) if meta.file_type().is_symlink() => { + targets.dirs.insert(resolved.clone()); + targets.entries.insert(entry_key(&candidate)); + hops += 1; + let Ok(link) = std::fs::read_link(&candidate) else { + return; + }; + if hops > MAX_SYMLINK_HOPS { + warn!("Too many levels of symlinks while resolving {:?}", path); + return; + } + // Relative targets resolve against `resolved`; absolute + // ones reset it through their root component. + for c in link.components().rev() { + pending.push_front(PathBuf::from(c.as_os_str())); + } + } + Ok(_) if !is_last => resolved = candidate, + Ok(_) | Err(_) => { + if is_last { + targets.file_dirs.insert(resolved.clone()); + } + targets.dirs.insert(resolved); + targets.entries.insert(entry_key(&candidate)); + return; + } + } + } + } + } +} + +/// Re-resolves `file_paths` and updates the watched directories, returning +/// the new targets. Failures are logged and retried on the next event that +/// matches a chain entry. +fn rearm_watches( + watcher: &mut impl Watcher, + file_paths: &[PathBuf], + old: WatchTargets, +) -> WatchTargets { + let mut new = resolve_watch_targets(file_paths); + + for dir in old.dirs.difference(&new.dirs) { + // The directory may already be gone; nothing to do then. + let _ = watcher.unwatch(dir); + } + let mut failed = Vec::new(); + for dir in new.dirs.difference(&old.dirs) { + match watcher.watch(dir, RecursiveMode::NonRecursive) { + Ok(()) => { + info!("File watcher now watching {:?}", dir); + warn_if_root(dir); + } + Err(err) => { + warn!("Failed to watch {:?}: {}", dir, err); + failed.push(dir.clone()); + } + } + } + for dir in failed { + new.dirs.remove(&dir); + } + new +} + +/// Warns when the root directory is watched, which happens when a path goes +/// through a symlink directly under it (e.g. `/tmp`, `/var` or `/etc` on +/// macOS). On macOS the FSEvents backend then receives every event on the +/// volume and filters them in user space. +fn warn_if_root(dir: &Path) { + if dir.parent().is_none() { + warn!( + "File watcher is watching the root directory {:?} because a watched path goes \ + through a symlink in it; on macOS this receives every event on the volume", + dir + ); + } +} + /// Check if an event kind is relevant based on the configuration. fn is_relevant_event(kind: &EventKind, config: &FileWatcherConfig) -> bool { match kind { @@ -189,47 +380,75 @@ fn is_relevant_event(kind: &EventKind, config: &FileWatcherConfig) -> bool { #[cfg(test)] mod tests { - use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; + use std::sync::{Arc, Mutex}; + use std::time::{Duration, Instant}; use common_test_util::temp_dir::create_temp_dir; use super::*; + const TIMEOUT: Duration = Duration::from_secs(10); + + /// Polls `cond` until it holds or `TIMEOUT` elapses. + fn wait_until(cond: impl Fn() -> bool) -> bool { + let deadline = Instant::now() + TIMEOUT; + while Instant::now() < deadline { + if cond() { + return true; + } + std::thread::sleep(Duration::from_millis(20)); + } + cond() + } + + /// Spawns a watcher on `path` and returns a counter of callback calls. + fn spawn_counter(path: &Path) -> Arc { + let counter = Arc::new(AtomicUsize::new(0)); + let counter_clone = counter.clone(); + FileWatcherBuilder::new() + .watch_path(path) + .unwrap() + .config(FileWatcherConfig::new()) + .spawn(move || { + counter_clone.fetch_add(1, Ordering::SeqCst); + }) + .unwrap(); + counter + } + + /// Spawns a watcher on `path` whose callback stores the content read + /// through `path`, like the real reload callbacks do. + fn spawn_reader(path: &Path) -> Arc> { + let content = Arc::new(Mutex::new(String::new())); + let content_clone = content.clone(); + let path_owned = path.to_path_buf(); + FileWatcherBuilder::new() + .watch_path(path) + .unwrap() + .spawn(move || { + if let Ok(c) = std::fs::read_to_string(&path_owned) { + *content_clone.lock().unwrap() = c; + } + }) + .unwrap(); + content + } + #[test] fn test_file_watcher_detects_changes() { common_telemetry::init_default_ut_logging(); let dir = create_temp_dir("test_file_watcher"); let file_path = dir.path().join("test_file.txt"); - - // Create initial file std::fs::write(&file_path, "initial content").unwrap(); - let counter = Arc::new(AtomicUsize::new(0)); - let counter_clone = counter.clone(); - - FileWatcherBuilder::new() - .watch_path(&file_path) - .unwrap() - .config(FileWatcherConfig::new()) - .spawn(move || { - counter_clone.fetch_add(1, Ordering::SeqCst); - }) - .unwrap(); - - // Give watcher time to start - std::thread::sleep(Duration::from_millis(100)); - - // Modify the file + // Watches are registered before `spawn` returns, so no need to wait. + let counter = spawn_counter(&file_path); std::fs::write(&file_path, "modified content").unwrap(); - // Wait for the event to be processed - std::thread::sleep(Duration::from_millis(500)); - assert!( - counter.load(Ordering::SeqCst) >= 1, + wait_until(|| counter.load(Ordering::SeqCst) >= 1), "Watcher should have detected at least one change" ); } @@ -240,38 +459,181 @@ mod tests { let dir = create_temp_dir("test_file_watcher_recreate"); let file_path = dir.path().join("test_file.txt"); - - // Create initial file std::fs::write(&file_path, "initial content").unwrap(); - let counter = Arc::new(AtomicUsize::new(0)); - let counter_clone = counter.clone(); + let content = spawn_reader(&file_path); - FileWatcherBuilder::new() - .watch_path(&file_path) - .unwrap() - .config(FileWatcherConfig::new()) - .spawn(move || { - counter_clone.fetch_add(1, Ordering::SeqCst); - }) - .unwrap(); - - // Give watcher time to start - std::thread::sleep(Duration::from_millis(100)); - - // Delete the file std::fs::remove_file(&file_path).unwrap(); - std::thread::sleep(Duration::from_millis(100)); - // Recreate the file - this should still be detected because we watch the directory std::fs::write(&file_path, "recreated content").unwrap(); - // Wait for the event to be processed - std::thread::sleep(Duration::from_millis(500)); - assert!( - counter.load(Ordering::SeqCst) >= 1, + wait_until(|| *content.lock().unwrap() == "recreated content"), "Watcher should have detected file recreation" ); } + + #[test] + fn test_file_watcher_ignores_unrelated_files() { + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_watcher_unrelated"); + let file_path = dir.path().join("watched.txt"); + std::fs::write(&file_path, "initial content").unwrap(); + + let counter = spawn_counter(&file_path); + + std::fs::write(dir.path().join("unrelated.txt"), "noise").unwrap(); + // Write the watched file afterwards: once its event is seen, the + // earlier unrelated event has been processed too. + std::fs::write(&file_path, "modified content").unwrap(); + assert!(wait_until(|| counter.load(Ordering::SeqCst) >= 1)); + std::thread::sleep(Duration::from_millis(200)); + let seen = counter.load(Ordering::SeqCst); + + // Only unrelated activity from now on. + std::fs::write(dir.path().join("unrelated.txt"), "more noise").unwrap(); + std::fs::remove_file(dir.path().join("unrelated.txt")).unwrap(); + std::thread::sleep(Duration::from_millis(500)); + assert_eq!( + counter.load(Ordering::SeqCst), + seen, + "Unrelated files in the watched directory must not trigger the callback" + ); + } + + #[cfg(unix)] + #[test] + fn test_resolve_watch_targets_symlink_chain() { + use std::os::unix::fs::symlink; + + let dir = create_temp_dir("test_resolve_watch_targets"); + let base = dir.path().canonicalize().unwrap(); + let keys = |paths: &[PathBuf]| paths.iter().map(|p| entry_key(p)).collect::>(); + + // real/users <- secrets (relative dir symlink) <- conf/users (relative + // file symlink going through `..`). + std::fs::create_dir(base.join("real")).unwrap(); + std::fs::write(base.join("real/users"), "").unwrap(); + symlink("real", base.join("secrets")).unwrap(); + std::fs::create_dir(base.join("conf")).unwrap(); + symlink("../secrets/users", base.join("conf/users")).unwrap(); + + let targets = resolve_watch_targets(&[base.join("conf/users")]); + assert_eq!( + targets.dirs, + HashSet::from([base.clone(), base.join("conf"), base.join("real")]) + ); + assert_eq!(targets.file_dirs, HashSet::from([base.join("real")])); + assert_eq!( + targets.entries, + keys(&[ + base.join("conf/users"), + base.join("secrets"), + base.join("real/users"), + ]) + ); + + // A plain file only needs its parent directory. + let targets = resolve_watch_targets(&[base.join("real/users")]); + assert_eq!(targets.dirs, HashSet::from([base.join("real")])); + assert_eq!(targets.entries, keys(&[base.join("real/users")])); + + // A dangling link keeps the directory where the target should appear. + std::fs::remove_file(base.join("secrets")).unwrap(); + let targets = resolve_watch_targets(&[base.join("conf/users")]); + assert_eq!( + targets.dirs, + HashSet::from([base.clone(), base.join("conf")]) + ); + assert!(targets.entries.contains(&entry_key(&base.join("secrets")))); + + // Symlink loops terminate. + symlink("loop", base.join("loop")).unwrap(); + let targets = resolve_watch_targets(&[base.join("loop")]); + assert_eq!(targets.dirs, HashSet::from([base.clone()])); + } + + #[cfg(unix)] + #[test] + fn test_file_watcher_detects_dir_symlink_swap() { + use std::os::unix::fs::symlink; + + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_watcher_symlink_swap"); + let base = dir.path(); + std::fs::create_dir(base.join("gen1")).unwrap(); + std::fs::write(base.join("gen1/file"), "1").unwrap(); + std::fs::create_dir(base.join("gen2")).unwrap(); + std::fs::write(base.join("gen2/file"), "2").unwrap(); + symlink(base.join("gen1"), base.join("current")).unwrap(); + std::fs::create_dir(base.join("conf")).unwrap(); + symlink(base.join("current/file"), base.join("conf/file")).unwrap(); + + let content = spawn_reader(&base.join("conf/file")); + + // Swap `current` atomically, like `ln -sfn gen2 current`. + symlink(base.join("gen2"), base.join("current.tmp")).unwrap(); + std::fs::rename(base.join("current.tmp"), base.join("current")).unwrap(); + assert!( + wait_until(|| *content.lock().unwrap() == "2"), + "Watcher should have detected the swap" + ); + // Let any remaining events of the swap drain, so they can't read the + // edit below by accident. + std::thread::sleep(Duration::from_millis(200)); + + // The watch follows the chain to the new target directory. + std::fs::write(base.join("gen2/file"), "22").unwrap(); + assert!( + wait_until(|| *content.lock().unwrap() == "22"), + "Watcher should have detected the edit of the new target" + ); + } + + /// A directory holding a symlink of the chain may be traversable but not + /// readable. It can't be watched, but that must not prevent watching the + /// file itself. + #[cfg(unix)] + #[test] + fn test_file_watcher_skips_unreadable_link_dir() { + use std::os::unix::fs::{PermissionsExt, symlink}; + + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_watcher_unreadable_dir"); + let base = dir.path(); + std::fs::create_dir(base.join("real")).unwrap(); + std::fs::write(base.join("real/file"), "1").unwrap(); + std::fs::create_dir(base.join("locked")).unwrap(); + symlink(base.join("real/file"), base.join("locked/file")).unwrap(); + + /// Restores a readable mode on drop, even if the test panics: the + /// temp dir can't be removed while `locked` is unreadable, since + /// removing it requires listing its entries. + struct RestoreMode(PathBuf); + impl Drop for RestoreMode { + fn drop(&mut self) { + let _ = std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(0o755)); + } + } + + // Declared after `dir`, so it drops (restoring the mode) first. + let _restore = RestoreMode(base.join("locked")); + // Traversable (so `locked/file` still resolves) but not readable. + std::fs::set_permissions(base.join("locked"), std::fs::Permissions::from_mode(0o311)) + .unwrap(); + if std::fs::read_dir(base.join("locked")).is_ok() { + // Running as root (or with CAP_DAC_OVERRIDE): nothing to test. + return; + } + + // Would panic if watching `locked` were fatal. + let content = spawn_reader(&base.join("locked/file")); + + // The directory of the final file is still watched. + std::fs::write(base.join("real/file"), "2").unwrap(); + assert!(wait_until(|| *content.lock().unwrap() == "2")); + } } diff --git a/src/common/datasource/Cargo.toml b/src/common/datasource/Cargo.toml index d1ca609c606..3953d144bad 100644 --- a/src/common/datasource/Cargo.toml +++ b/src/common/datasource/Cargo.toml @@ -39,6 +39,7 @@ parquet.workspace = true paste.workspace = true regex.workspace = true serde.workspace = true +serde_json.workspace = true snafu.workspace = true strum.workspace = true tokio.workspace = true diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index d016d4eabe9..d4cfb45088c 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -647,16 +647,15 @@ mod tests { let compressed_file_name = format!("test_compressed_csv.{}", compression_type.file_extension()); let compressed_file_path = temp_dir.path().join(&compressed_file_name); - let compressed_file_path_str = compressed_file_path.to_str().unwrap(); // Create a simple file store for testing - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); // Export CSV with compression let rows = stream_to_csv( Box::pin(DfRecordBatchStreamAdapter::new(recordbatches.as_stream())), store, - compressed_file_path_str, + &compressed_file_name, 1024, 1, &format, @@ -708,13 +707,13 @@ mod tests { } // Verify the compressed file can be decompressed and content matches original data - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new( CsvFormat { compression_type, ..Default::default() } - .infer_schema(&store, compressed_file_path_str) + .infer_schema(&store, &compressed_file_name) .await .unwrap(), ); @@ -722,7 +721,7 @@ mod tests { let stream = file_to_stream( &store, - compressed_file_path_str, + &compressed_file_name, csv_source.clone(), None, compression_type, @@ -761,13 +760,13 @@ mod tests { ) .unwrap(); - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new(arrow_schema::Schema::new(vec![ Field::new("id", DataType::UInt32, false), Field::new("name", DataType::Utf8, false), Field::new("value", DataType::Float64, false), ])); - let path = csv_file_path.to_str().unwrap(); + let path = "input.csv"; let stream = tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) @@ -793,13 +792,13 @@ mod tests { let csv_file_path = temp_dir.path().join("input.csv"); std::fs::write(&csv_file_path, "id,name,value\n1,Alice,10.5\n2,Bob\n").unwrap(); - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new(arrow_schema::Schema::new(vec![ Field::new("id", DataType::UInt32, false), Field::new("name", DataType::Utf8, false), Field::new("value", DataType::Float64, false), ])); - let path = csv_file_path.to_str().unwrap(); + let path = "input.csv"; let stream = tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) diff --git a/src/common/datasource/src/file_format/json.rs b/src/common/datasource/src/file_format/json.rs index 3c24a4b7e27..23ce461fcb2 100644 --- a/src/common/datasource/src/file_format/json.rs +++ b/src/common/datasource/src/file_format/json.rs @@ -313,16 +313,15 @@ mod tests { let compressed_file_name = format!("test_compressed_json.{}", compression_type.file_extension()); let compressed_file_path = temp_dir.path().join(&compressed_file_name); - let compressed_file_path_str = compressed_file_path.to_str().unwrap(); // Create a simple file store for testing - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); // Export JSON with compression let rows = stream_to_json( Box::pin(DfRecordBatchStreamAdapter::new(recordbatches.as_stream())), store, - compressed_file_path_str, + &compressed_file_name, 1024, 1, &format, @@ -374,13 +373,13 @@ mod tests { } // Verify the compressed file can be decompressed and content matches original data - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new( JsonFormat { compression_type, ..Default::default() } - .infer_schema(&store, compressed_file_path_str) + .infer_schema(&store, &compressed_file_name) .await .unwrap(), ); @@ -388,7 +387,7 @@ mod tests { let stream = file_to_stream( &store, - compressed_file_path_str, + &compressed_file_name, json_source.clone(), None, compression_type, diff --git a/src/common/datasource/src/file_format/tests.rs b/src/common/datasource/src/file_format/tests.rs index 27b86e1b114..14be2b1777f 100644 --- a/src/common/datasource/src/file_format/tests.rs +++ b/src/common/datasource/src/file_format/tests.rs @@ -66,13 +66,12 @@ impl Test<'_> { #[tokio::test] async fn test_json_opener() { - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = basic_schema_with_time_format(); let file_source = Arc::new(JsonSource::new(schema)).with_batch_size(test_util::TEST_BATCH_SIZE); - let path = &find_workspace_path("/src/common/datasource/tests/json/basic.json") - .display() - .to_string(); + let path = "json/basic.json"; let tests = [ Test { @@ -108,11 +107,10 @@ async fn test_json_opener() { #[tokio::test] async fn test_csv_opener() { - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = basic_schema_with_time_format(); - let path = &find_workspace_path("/src/common/datasource/tests/csv/basic.csv") - .display() - .to_string(); + let path = "csv/basic.csv"; let file_source = CsvSource::new(schema).with_batch_size(test_util::TEST_BATCH_SIZE); @@ -150,13 +148,12 @@ async fn test_csv_opener() { #[tokio::test(flavor = "multi_thread")] async fn test_parquet_exec() { - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = test_basic_schema(); - let path = &find_workspace_path("/src/common/datasource/tests/parquet/basic.parquet") - .display() - .to_string(); + let path = "parquet/basic.parquet"; let parquet_source = ParquetSource::new(schema) .with_parquet_file_reader_factory(Arc::new(DefaultParquetFileReaderFactory::new(store))); @@ -191,11 +188,10 @@ async fn test_parquet_exec() { #[tokio::test] async fn test_orc_opener() { - let path = &find_workspace_path("/src/common/datasource/tests/orc/test.orc") - .display() - .to_string(); + let path = "orc/test.orc"; - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = Arc::new(OrcFormat.infer_schema(&store, path).await.unwrap()); let file_source = Arc::new(OrcSource::new(schema.into())); diff --git a/src/common/datasource/src/lib.rs b/src/common/datasource/src/lib.rs index c832831da80..ebec28d7129 100644 --- a/src/common/datasource/src/lib.rs +++ b/src/common/datasource/src/lib.rs @@ -20,6 +20,7 @@ pub mod file_format; pub mod lister; pub mod object_store; pub mod packed_snapshot; +pub mod packed_writer; pub mod parquet_writer; pub mod share_buffer; #[cfg(test)] diff --git a/src/common/datasource/src/packed_writer.rs b/src/common/datasource/src/packed_writer.rs new file mode 100644 index 00000000000..635a1a7a771 --- /dev/null +++ b/src/common/datasource/src/packed_writer.rs @@ -0,0 +1,669 @@ +// 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. + +//! Streaming destinations for one schema chunk of independent Parquet streams. + +use std::sync::Arc; + +use bytes::Bytes; +use object_store::{ObjectStore, Writer}; +use snafu::{ResultExt, ensure}; +use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, Result}; +use crate::packed_snapshot::{ObjectKind, PACK_INDEX_FILE, PackIndex, PackObject, PackTable}; + +pub const WRITE_BYTES: usize = 8 * 1024 * 1024; + +/// One request owns this sink. Finishers retain their encoder admission while +/// waiting on the mutex, bounding both completed buffers and pending appends. +pub type PackedWriterRef = Arc>; + +#[derive(Clone, Copy)] +struct UploadLimits { + part: usize, + object: u64, +} + +impl UploadLimits { + fn for_store(store: &ObjectStore) -> Result { + let info = store.info(); + let caps = info.capability(); + let part = caps + .write_multi_max_size + .unwrap_or(WRITE_BYTES) + .min(WRITE_BYTES); + ensure!( + part > 0 && caps.write_multi_min_size.unwrap_or(0) <= part, + error::ParquetWriterResourceSnafu { + reason: "backend cannot accept bounded 8 MiB parts" + } + ); + // OpenDAL exposes part size but not part count. These are the limits of + // its multipart implementations, including GCS's XML multipart writer. + let parts = match info.scheme() { + "s3" | "oss" | "gcs" => 10_000, + "azblob" => 50_000, + _ => u64::MAX, + }; + let object = (caps.write_total_max_size.unwrap_or(usize::MAX) as u64) + .min(parts.saturating_mul(part as u64)); + Ok(Self { part, object }) + } +} + +struct Upload { + writer: Writer, + store: ObjectStore, + path: String, + limits: UploadLimits, + length: u64, + conditional: bool, + close_started: bool, +} + +impl Upload { + async fn open(store: &ObjectStore, path: String, limits: UploadLimits) -> Result { + let conditional = store.info().capability().write_with_if_not_exists; + if !conditional { + ensure!( + !store + .exists(&path) + .await + .context(error::ReadObjectSnafu { path: &path })?, + error::InvalidPackedSnapshotSnafu { + reason: format!("output already exists: {path}") + } + ); + } + let writer = store + .writer_with(&path) + .chunk(limits.part) + .concurrent(1) + .if_not_exists(conditional) + .await + .context(error::WriteObjectSnafu { path: &path })?; + Ok(Self { + writer, + store: store.clone(), + path, + limits, + length: 0, + conditional, + close_started: false, + }) + } + + async fn write(&mut self, bytes: Bytes) -> Result<()> { + ensure!( + (bytes.len() as u64) <= self.limits.object.saturating_sub(self.length), + error::ParquetWriterResourceSnafu { + reason: "packed export object or multipart part-count limit exceeded" + } + ); + for offset in (0..bytes.len()).step_by(self.limits.part) { + self.writer + .write(bytes.slice(offset..(offset + self.limits.part).min(bytes.len()))) + .await + .context(error::WriteObjectSnafu { path: &self.path })?; + } + self.length += bytes.len() as u64; + Ok(()) + } + + async fn write_json_array( + &mut self, + values: &[T], + token: &CancellationToken, + ) -> Result<()> { + self.write(Bytes::from_static(b"[")).await?; + for (i, value) in values.iter().enumerate() { + check_cancelled(Some(token))?; + if i > 0 { + self.write(Bytes::from_static(b",")).await?; + } + let bytes = serde_json::to_vec(value).map_err(|e| { + error::InvalidPackedSnapshotSnafu { + reason: e.to_string(), + } + .build() + })?; + self.write(bytes.into()).await?; + } + self.write(Bytes::from_static(b"]")).await + } + + async fn close(&mut self) -> Result<()> { + self.close_started = true; + self.writer + .close() + .await + .context(error::WriteObjectSnafu { path: &self.path })?; + Ok(()) + } + + async fn abort(mut self) -> Result<()> { + let result = self.writer.abort().await; + if !self.conditional + && result.as_ref().is_err_and(|e| { + e.kind() == object_store::ErrorKind::Unsupported + && (!self.close_started + || object_store::secure_fs::is_unsynced_overwrite_abort(e)) + }) + { + let store = self.store.clone(); + let path = self.path.clone(); + drop(self); + store + .delete(&path) + .await + .context(error::WriteObjectSnafu { path })?; + } else { + result.context(error::WriteObjectSnafu { path: &self.path })?; + } + Ok(()) + } +} + +/// Pack/index metadata is accounted separately from retained Arrow payloads. +pub struct PackedWriter { + store: ObjectStore, + limits: UploadLimits, + pack: Option, + index: PackIndex, + next_pack: usize, +} + +impl PackedWriter { + pub fn new(store: ObjectStore) -> Result { + let limits = UploadLimits::for_store(&store)?; + Ok(Arc::new(Mutex::new(Self { + store, + limits, + pack: None, + index: PackIndex { + version: 1, + objects: vec![], + tables: vec![], + }, + next_pack: 0, + }))) + } + + async fn close_pack(&mut self) -> Result<()> { + if let Some(pack) = &mut self.pack { + pack.close().await?; + self.index.objects.push(PackObject { + path: pack.path.clone(), + kind: ObjectKind::Pack, + length: pack.length, + }); + self.pack = None; + } + Ok(()) + } + + async fn append( + &mut self, + name: String, + bytes: Bytes, + rows: u64, + token: Option<&CancellationToken>, + ) -> Result<()> { + check_cancelled(token)?; + if self + .pack + .as_ref() + .is_some_and(|p| bytes.len() as u64 > self.limits.object.saturating_sub(p.length)) + { + self.close_pack().await?; + } + check_cancelled(token)?; + if self.pack.is_none() { + let path = format!("pack-{:06}.bin", self.next_pack); + self.next_pack += 1; + self.pack = Some(Upload::open(&self.store, path, self.limits).await?); + } + let pack = self.pack.as_mut().ok_or_else(|| { + error::InvalidPackedSnapshotSnafu { + reason: "missing pack writer", + } + .build() + })?; + let table = PackTable { + table_name: name, + object: pack.path.clone(), + offset: pack.length, + length: bytes.len() as u64, + row_count: rows, + }; + pack.write(bytes).await?; + check_cancelled(token)?; + self.index.tables.push(table); + Ok(()) + } + + /// Publish the index only after every data object has closed successfully. + pub async fn finish(&mut self, token: &CancellationToken) -> Result> { + check_cancelled(Some(token))?; + self.close_pack().await?; + check_cancelled(Some(token))?; + self.index + .tables + .sort_unstable_by(|a, b| a.table_name.cmp(&b.table_name)); + self.index.validate()?; + let mut upload = Upload::open(&self.store, PACK_INDEX_FILE.into(), self.limits).await?; + let result = async { + upload + .write(Bytes::from_static(b"{\"version\":1,\"objects\":")) + .await?; + upload.write_json_array(&self.index.objects, token).await?; + upload.write(Bytes::from_static(b",\"tables\":")).await?; + upload.write_json_array(&self.index.tables, token).await?; + upload.write(Bytes::from_static(b"}")).await?; + check_cancelled(Some(token))?; + upload.close().await?; + check_cancelled(Some(token)) + } + .await; + if let Err(error) = result { + if let Err(secondary) = upload.abort().await { + common_telemetry::warn!(secondary; "Failed to abort pack index"); + } + return Err(error); + } + Ok(self + .index + .objects + .iter() + .map(|o| o.path.clone()) + .chain([PACK_INDEX_FILE.into()]) + .collect()) + } + + /// Call only after all admitted encoders have drained their I/O. + pub async fn abort(&mut self) -> Result<()> { + if let Some(pack) = self.pack.take() { + pack.abort().await?; + } + Ok(()) + } +} + +/// Buffers a small stream, or spills its prefix and continues the same stream. +pub struct PackedTableWriter { + shared: PackedWriterRef, + name: String, + path: String, + buffer: Vec, + ordinary: bool, + standalone: Option, +} + +impl PackedTableWriter { + /// The caller retains encoder admission until finish/abort returns. + pub fn new(shared: PackedWriterRef, name: String, id: u32, ordinary: bool) -> Self { + Self { + shared, + name, + path: format!("table-{id}.parquet"), + buffer: if ordinary { + Vec::new() + } else { + Vec::with_capacity(WRITE_BYTES) + }, + ordinary, + standalone: None, + } + } + + pub(crate) async fn write(&mut self, bytes: Bytes) -> Result<()> { + if self.standalone.is_none() + && (self.ordinary || self.buffer.len().saturating_add(bytes.len()) > WRITE_BYTES) + { + let sink = self.shared.lock().await; + self.standalone = + Some(Upload::open(&sink.store, self.path.clone(), sink.limits).await?); + } + if let Some(upload) = &mut self.standalone { + if !self.buffer.is_empty() { + upload + .write(std::mem::take(&mut self.buffer).into()) + .await?; + } + upload.write(bytes).await + } else { + self.buffer.extend_from_slice(&bytes); + Ok(()) + } + } + + pub(crate) async fn finish( + &mut self, + rows: u64, + token: Option<&CancellationToken>, + ) -> Result<()> { + check_cancelled(token)?; + if let Some(upload) = &mut self.standalone { + upload.close().await?; + check_cancelled(token)?; + let mut shared = self.shared.lock().await; + shared.index.objects.push(PackObject { + path: self.path.clone(), + kind: ObjectKind::Parquet, + length: upload.length, + }); + shared.index.tables.push(PackTable { + table_name: self.name.clone(), + object: self.path.clone(), + offset: 0, + length: upload.length, + row_count: rows, + }); + } else { + self.shared + .lock() + .await + .append( + self.name.clone(), + // OpenDAL may retain each slice until the pack part fills. + // Release the encoder's 8 MiB capacity before queuing it. + Bytes::from(std::mem::take(&mut self.buffer).into_boxed_slice()), + rows, + token, + ) + .await?; + } + Ok(()) + } + + pub(crate) async fn abort(&mut self) -> Result<()> { + if let Some(upload) = self.standalone.take() { + upload.abort().await?; + } + Ok(()) + } +} + +fn check_cancelled(token: Option<&CancellationToken>) -> Result<()> { + ensure!( + token.is_none_or(|t| !t.is_cancelled()), + error::ParquetWriteCancelledSnafu + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use arrow::array::{ArrayRef, Int64Array, StringArray}; + use arrow::record_batch::RecordBatch; + use futures::TryStreamExt; + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + use parquet::arrow::ParquetRecordBatchStreamBuilder; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + use super::*; + use crate::file_format::parquet::packed_reader::{PackReadWindows, PackedParquetReader}; + use crate::parquet_writer::{ParquetFileWriter, ParquetWriterLimits}; + + #[tokio::test] + async fn independent_streams_empty_and_spilled_roundtrip() { + let directory = common_test_util::temp_dir::create_temp_dir("packed-writer"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let factory: MockWriterFactory = Arc::new({ + let sizes = sizes.clone(); + move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let shared = PackedWriter::new(store.clone()).unwrap(); + let small = RecordBatch::try_from_iter([( + "tag", + Arc::new(StringArray::from(vec![Some(""), None, Some("value")])) as ArrayRef, + )]) + .unwrap(); + let mut state = 17u64; + let values = (0..600_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as i64 + }) + .collect::>(); + let random = Arc::new(Int64Array::from(values)) as ArrayRef; + let large = RecordBatch::try_from_iter([("a", random.clone()), ("b", random)]).unwrap(); + let batches = [small.clone(), RecordBatch::new_empty(large.schema()), large]; + for (id, batch) in batches.iter().enumerate() { + let sink = + PackedTableWriter::new(shared.clone(), format!("table{id}"), id as u32, false); + let mut writer = ParquetFileWriter::open_packed( + batch.schema(), + store.clone(), + "unused", + Some(ParquetWriterLimits { + row_group_rows: 8192, + flush_threshold_bytes: WRITE_BYTES, + max_row_groups: 4096, + }), + sink, + ) + .unwrap(); + writer.write(batch.clone(), None).await.unwrap(); + writer.finish(None).await.unwrap(); + } + let inventory = shared + .lock() + .await + .finish(&CancellationToken::new()) + .await + .unwrap(); + let index: PackIndex = + serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap(); + index + .validate_membership(["table0", "table1", "table2"]) + .unwrap(); + assert_eq!(inventory.len(), 3); + assert!(sizes.lock().unwrap().iter().all(|n| *n <= WRITE_BYTES)); + assert_eq!( + index + .objects + .iter() + .filter(|o| o.kind == ObjectKind::Pack) + .count(), + 1 + ); + assert!( + index + .objects + .iter() + .any(|o| o.kind == ObjectKind::Parquet && o.length > WRITE_BYTES as u64) + ); + let windows = PackReadWindows::new(store.clone()); + for (entry, expected) in index.tables.iter().zip(batches) { + let object = index + .objects + .iter() + .find(|o| o.path == entry.object) + .unwrap(); + let actual = if object.kind == ObjectKind::Pack { + let reader = PackedParquetReader::new( + windows.clone(), + object.path.clone(), + object.length, + entry.offset, + entry.length, + ) + .unwrap(); + ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .build() + .unwrap() + .try_collect::>() + .await + .unwrap() + } else { + ParquetRecordBatchReaderBuilder::try_new( + store.read(&object.path).await.unwrap().to_bytes(), + ) + .unwrap() + .build() + .unwrap() + .collect::, _>>() + .unwrap() + }; + assert_eq!(entry.row_count, expected.num_rows() as u64); + assert_eq!( + arrow::compute::concat_batches(&expected.schema(), &actual).unwrap(), + expected + ); + } + } + + struct Observe(oio::Writer, Arc>>); + impl oio::Write for Observe { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.1.lock().unwrap().push(bytes.len()); + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + + #[tokio::test] + async fn rolls_at_table_boundary_and_aborts_standalone_at_part_limit() { + let directory = common_test_util::temp_dir::create_temp_dir("packed-limit"); + let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let factory: MockWriterFactory = Arc::new({ + let sizes = sizes.clone(); + move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + }); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let shared = PackedWriter::new(store.clone()).unwrap(); + // Two data parts per object; no oversized fixture is needed. + shared.lock().await.limits = UploadLimits { + part: 512, + object: 1024, + }; + for id in 0..3 { + let mut table = PackedTableWriter::new(shared.clone(), format!("t{id}"), id, false); + table.write(vec![id as u8; 600].into()).await.unwrap(); + table.finish(0, None).await.unwrap(); + } + let mut table = PackedTableWriter::new(shared.clone(), "large".into(), 4, true); + table.write(vec![0; 1024].into()).await.unwrap(); + assert!(matches!( + table.write(Bytes::from_static(b"x")).await, + Err(error::Error::ParquetWriterResource { .. }) + )); + table.abort().await.unwrap(); + assert!(!store.exists("table-4.parquet").await.unwrap()); + shared + .lock() + .await + .finish(&CancellationToken::new()) + .await + .unwrap(); + let index: PackIndex = + serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap(); + assert_eq!(index.objects.len(), 3); + assert!( + index + .tables + .iter() + .all(|t| t.offset == 0 && t.length == 600) + ); + assert!(sizes.lock().unwrap().iter().all(|n| *n <= 512)); + } + + struct FailClose(oio::Writer); + impl oio::Write for FailClose { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected close failure", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + + #[tokio::test] + async fn close_failure_cancellation_and_collision_never_publish_index() { + for failure in ["pack-000000.bin", PACK_INDEX_FILE, "cancel", "collision"] { + let directory = common_test_util::temp_dir::create_temp_dir("packed-failure"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + if failure == "collision" { + store.write("pack-000000.bin", "winner").await.unwrap(); + } + let factory: MockWriterFactory = Arc::new(move |path, _, writer| { + if path.trim_start_matches('/') == failure { + Box::new(FailClose(writer)) + } else { + writer + } + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let shared = PackedWriter::new(store.clone()).unwrap(); + let mut table = PackedTableWriter::new(shared.clone(), "t".into(), 0, false); + table.write(vec![0; 12].into()).await.unwrap(); + table.finish(0, None).await.unwrap(); + let token = CancellationToken::new(); + if failure == "cancel" { + token.cancel(); + } + assert!(shared.lock().await.finish(&token).await.is_err()); + let _ = shared.lock().await.abort().await; + assert!(!store.exists(PACK_INDEX_FILE).await.unwrap()); + if failure == "collision" { + assert_eq!( + store.read("pack-000000.bin").await.unwrap().to_bytes(), + Bytes::from_static(b"winner") + ); + } + } + } +} diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index ae1fcec04a9..243a46b4e1a 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -28,6 +28,14 @@ use tokio_util::sync::CancellationToken; use crate::DEFAULT_WRITE_BUFFER_SIZE; use crate::error::{self, Result}; +use crate::packed_writer::PackedTableWriter; + +/// Destination creation policy; conditional failures never authorize path deletion. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ParquetCreationPolicy { + Overwrite, + IfNotExists, +} /// Limits for one Parquet file. Flush thresholds are not hard memory caps. #[derive(Clone, Copy, Debug)] @@ -45,10 +53,18 @@ pub struct ParquetWriterLimits { /// operation can leave a filesystem worker running after cleanup. pub struct ParquetFileWriter { encoder: Option>>, - sink: Writer, + sink: ParquetSink, + rows: u64, store: ObjectStore, path: String, limits: Option, + creation: ParquetCreationPolicy, + close_started: bool, +} + +enum ParquetSink { + Object(Writer), + Packed(Box), } impl ParquetFileWriter { @@ -61,43 +77,63 @@ impl ParquetFileWriter { concurrency: usize, limits: Option, ) -> Result { - let mut props = WriterProperties::builder() - .set_compression(Compression::ZSTD(ZstdLevel::default())) - .set_statistics_truncate_length(None) - .set_column_index_truncate_length(None); - if let Some(limits) = limits { - ensure!( - limits.row_group_rows > 0 - && limits.flush_threshold_bytes > 0 - && limits.max_row_groups > 0, - error::InvalidParquetWriterLimitsSnafu - ); - props = props - .set_max_row_group_row_count(Some(limits.row_group_rows)) - .set_max_row_group_bytes(None); - } - for field in schema.fields() { - if matches!(field.data_type(), DataType::Timestamp(_, _)) { - let column = ColumnPath::new(vec![field.name().clone()]); - props = props - .set_column_dictionary_enabled(column.clone(), false) - .set_column_encoding(column, Encoding::DELTA_BINARY_PACKED); - } - } - let encoder = ArrowWriter::try_new(Vec::new(), schema, Some(props.build())) - .context(error::WriteParquetSnafu { path })?; + Self::open_with_creation( + schema, + store, + path, + concurrency, + limits, + ParquetCreationPolicy::Overwrite, + ) + .await + } + + /// Open with a conditional policy only when the caller verified backend support. + pub async fn open_with_creation( + schema: SchemaRef, + store: ObjectStore, + path: &str, + concurrency: usize, + limits: Option, + creation: ParquetCreationPolicy, + ) -> Result { + let encoder = build_encoder(schema, limits, path)?; let sink = store .writer_with(path) .concurrent(concurrency) .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .if_not_exists(creation == ParquetCreationPolicy::IfNotExists) .await .context(error::WriteObjectSnafu { path })?; Ok(Self { encoder: Some(encoder), - sink, + sink: ParquetSink::Object(sink), + rows: 0, store, path: path.to_owned(), limits, + creation, + close_started: false, + }) + } + + /// Encode into a request-owned pack or standalone destination. + pub fn open_packed( + schema: SchemaRef, + store: ObjectStore, + path: &str, + limits: Option, + packed: PackedTableWriter, + ) -> Result { + Ok(Self { + encoder: Some(build_encoder(schema, limits, path)?), + sink: ParquetSink::Packed(Box::new(packed)), + rows: 0, + store, + path: path.into(), + limits, + creation: ParquetCreationPolicy::IfNotExists, + close_started: false, }) } @@ -150,15 +186,22 @@ impl ParquetFileWriter { check_cancelled(cancellation)?; self.write_bytes(bytes).await?; check_cancelled(cancellation)?; + self.rows += len as u64; offset += len; } Ok(()) } async fn write_bytes(&mut self, bytes: Vec) -> Result<()> { - if !bytes.is_empty() { - self.sink - .write(bytes) + let bytes = Bytes::from(bytes); + let sink = match &mut self.sink { + ParquetSink::Packed(packed) => return packed.write(bytes).await, + ParquetSink::Object(sink) => sink, + }; + let chunk = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; + // Slices retain the complete encoded allocation until its last submission. + for offset in (0..bytes.len()).step_by(chunk) { + sink.write(bytes.slice(offset..(offset + chunk).min(bytes.len()))) .await .context(error::WriteObjectSnafu { path: &self.path })?; } @@ -183,21 +226,31 @@ impl ParquetFileWriter { .context(error::JoinHandleSnafu)??; self.write_bytes(bytes).await?; check_cancelled(cancellation)?; - self.sink - .close() + let sink = match &mut self.sink { + ParquetSink::Packed(packed) => return packed.finish(self.rows, cancellation).await, + ParquetSink::Object(sink) => sink, + }; + self.close_started = true; + sink.close() .await .context(error::WriteObjectSnafu { path: &self.path })?; check_cancelled(cancellation)?; Ok(()) } - /// Abort an exclusively owned file after all in-flight operations have completed. - /// If the backend cannot abort, deletion assumes this attempt owns the path. + /// Abort after all in-flight operations complete. Preserve ambiguous commits; + /// conditional callers delegate cleanup exclusively to the backend. pub async fn abort(mut self) -> Result<()> { - let result = self.sink.abort().await; - if result - .as_ref() - .is_err_and(|e| e.kind() == object_store::ErrorKind::Unsupported) + let result = match &mut self.sink { + ParquetSink::Packed(packed) => return packed.abort().await, + ParquetSink::Object(sink) => sink.abort().await, + }; + if self.creation == ParquetCreationPolicy::Overwrite + && result.as_ref().is_err_and(|error| { + error.kind() == object_store::ErrorKind::Unsupported + && (!self.close_started + || object_store::secure_fs::is_unsynced_overwrite_abort(error)) + }) { let store = self.store.clone(); let path = self.path.clone(); @@ -214,6 +267,39 @@ impl ParquetFileWriter { } } +fn build_encoder( + schema: SchemaRef, + limits: Option, + path: &str, +) -> Result>> { + let mut props = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_statistics_truncate_length(None) + .set_column_index_truncate_length(None); + if let Some(limits) = limits { + ensure!( + limits.row_group_rows > 0 + && limits.flush_threshold_bytes > 0 + && limits.max_row_groups > 0, + error::InvalidParquetWriterLimitsSnafu + ); + props = props + .set_max_row_group_row_count(Some(limits.row_group_rows)) + .set_max_row_group_bytes(None); + } + for field in schema.fields() { + if matches!(field.data_type(), DataType::Timestamp(_, _)) { + let column = ColumnPath::new(vec![field.name().clone()]); + props = props + .set_column_dictionary_enabled(column.clone(), false) + .set_column_encoding(column, Encoding::DELTA_BINARY_PACKED); + } + } + let encoder = ArrowWriter::try_new(Vec::new(), schema, Some(props.build())) + .context(error::WriteParquetSnafu { path })?; + Ok(encoder) +} + fn check_cancelled(cancellation: Option<&CancellationToken>) -> Result<()> { ensure!( cancellation.is_none_or(|token| !token.is_cancelled()), @@ -288,6 +374,196 @@ mod tests { .unwrap() } + #[tokio::test] + async fn abort_preserves_collisions_and_ambiguous_commits() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio}; + + struct AmbiguousCommit(oio::Writer); + impl oio::Write for AmbiguousCommit { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close( + &mut self, + ) -> object_store::Result { + self.0.close().await?; + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "lost close reply", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + Err(object_store::Error::new( + object_store::ErrorKind::Unsupported, + "cannot abort", + )) + } + } + + for (creation, existing) in [ + (ParquetCreationPolicy::Overwrite, false), + (ParquetCreationPolicy::IfNotExists, false), + (ParquetCreationPolicy::IfNotExists, true), + ] { + let directory = common_test_util::temp_dir::create_temp_dir("conditional_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let path = "conditional.parquet"; + if existing { + store.write(path, "original").await.unwrap(); + } + let factory: MockWriterFactory = + Arc::new(|_, _, writer| Box::new(AmbiguousCommit(writer))); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut writer = ParquetFileWriter::open_with_creation( + batch().schema(), + store.clone(), + path, + 1, + None, + creation, + ) + .await + .unwrap(); + writer.write(batch(), None).await.unwrap(); + assert!(writer.finish(None).await.is_err()); + assert!(writer.abort().await.is_err()); + if existing { + assert_eq!( + store.read(path).await.unwrap().to_bytes(), + Bytes::from_static(b"original") + ); + } else { + assert_eq!( + read(&store, path) + .await + .metadata() + .file_metadata() + .num_rows(), + 4 + ); + } + } + } + + #[tokio::test] + async fn overwrite_abort_deletes_after_unsynced_close() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio}; + + struct FailedClose(oio::Writer); + impl oio::Write for FailedClose { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close( + &mut self, + ) -> object_store::Result { + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "close failed before sync", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + + let directory = common_test_util::temp_dir::create_temp_dir("unsynced_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let path = "partial.parquet"; + let factory: MockWriterFactory = Arc::new(|_, _, writer| Box::new(FailedClose(writer))); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut writer = ParquetFileWriter::open(batch().schema(), store.clone(), path, 1, None) + .await + .unwrap(); + writer.write(batch(), None).await.unwrap(); + assert!(writer.finish(None).await.is_err()); + assert!(store.exists(path).await.unwrap()); + writer.abort().await.unwrap(); + assert!(!store.exists(path).await.unwrap()); + } + + #[tokio::test] + async fn large_footer_flush_uses_bounded_submissions_in_one_parquet_stream() { + use std::sync::Mutex; + + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + struct ObservedWriter(oio::Writer, Arc>>); + impl oio::Write for ObservedWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.1.lock().unwrap().push(bytes.len()); + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + let directory = common_test_util::temp_dir::create_temp_dir("bounded_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let sizes = Arc::new(Mutex::new(Vec::new())); + let factory: MockWriterFactory = Arc::new({ + let sizes = sizes.clone(); + move |_, _, writer| Box::new(ObservedWriter(writer, sizes.clone())) + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut state = 17u64; + let values = (0..600_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as i64 + }) + .collect::>(); + let array = Arc::new(Int64Array::from(values)) as arrow::array::ArrayRef; + let batch = RecordBatch::try_from_iter([("a", array.clone()), ("b", array)]).unwrap(); + let mut writer = + ParquetFileWriter::open(batch.schema(), store.clone(), "large.parquet", 1, None) + .await + .unwrap(); + // No storage-layer chunking: observe the application's actual submissions. + writer.sink = ParquetSink::Object(store.writer("large.parquet").await.unwrap()); + writer.write(batch.clone(), None).await.unwrap(); + writer.finish(None).await.unwrap(); + let sizes = sizes.lock().unwrap().clone(); + let limit = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; + assert!(sizes.iter().sum::() > limit); + assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}"); + let actual = read(&store, "large.parquet") + .await + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!( + arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(), + batch + ); + } + #[tokio::test] async fn row_and_byte_limits_split_batches() { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); @@ -472,7 +748,8 @@ mod tests { .await .unwrap(); // Force footer bytes through the sink before its close operation. - writer.sink = store.writer_with("cancel.parquet").chunk(1).await.unwrap(); + writer.sink = + ParquetSink::Object(store.writer_with("cancel.parquet").chunk(1).await.unwrap()); let cancellation = CancellationToken::new(); let result = { let finish = writer.finish(Some(&cancellation)); diff --git a/src/common/datasource/src/test_util.rs b/src/common/datasource/src/test_util.rs index 4cefda88890..f95d24930b2 100644 --- a/src/common/datasource/src/test_util.rs +++ b/src/common/datasource/src/test_util.rs @@ -56,8 +56,8 @@ pub fn test_store(root: &str) -> ObjectStore { pub fn test_tmp_store(root: &str) -> (ObjectStore, TempDir) { let dir = create_temp_dir(root); - let builder = Fs::default(); - (ObjectStore::new(builder.root("/")).unwrap(), dir) + let store = test_store(dir.path().to_str().unwrap()); + (store, dir) } pub fn test_basic_schema() -> SchemaRef { @@ -94,9 +94,11 @@ pub(crate) fn scan_config( .build() } -pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usize) -> usize) { - let store = test_store("/"); - +pub async fn setup_stream_to_json_test( + store: &ObjectStore, + origin_path: &str, + threshold: impl Fn(usize) -> usize, +) { let schema = basic_schema_with_time_format(); let json_opener = JsonOpener::new( @@ -117,9 +119,9 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi .build() .unwrap(); - let (tmp_store, dir) = test_tmp_store("test_stream_to_json"); + let (tmp_store, _dir) = test_tmp_store("test_stream_to_json"); - let output_path = format!("{}/{}", dir.path().display(), "output"); + let output_path = "output"; let json_format = JsonFormat::default(); @@ -127,7 +129,7 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi stream_to_json( Box::pin(stream), tmp_store.clone(), - &output_path, + output_path, threshold(size), 8, &json_format, @@ -136,18 +138,17 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi .is_ok() ); - let written = tmp_store.read(&output_path).await.unwrap(); + let written = tmp_store.read(output_path).await.unwrap(); let origin = store.read(origin_path).await.unwrap(); assert_eq_lines(written.to_vec(), origin.to_vec()); } pub async fn setup_stream_to_csv_test( + store: &ObjectStore, origin_path: &str, format_path: &str, threshold: impl Fn(usize) -> usize, ) { - let store = test_store("/"); - let schema = basic_schema_with_time_format(); let csv_source = CsvSource::new(schema).with_batch_size(TEST_BATCH_SIZE); @@ -168,9 +169,9 @@ pub async fn setup_stream_to_csv_test( .build() .unwrap(); - let (tmp_store, dir) = test_tmp_store("test_stream_to_csv"); + let (tmp_store, _dir) = test_tmp_store("test_stream_to_csv"); - let output_path = format!("{}/{}", dir.path().display(), "output"); + let output_path = "output"; let csv_format = CsvFormat { timestamp_format: Some("%m-%d-%Y".to_string()), @@ -183,7 +184,7 @@ pub async fn setup_stream_to_csv_test( stream_to_csv( Box::pin(stream), tmp_store.clone(), - &output_path, + output_path, threshold(size), 8, &csv_format, @@ -192,7 +193,7 @@ pub async fn setup_stream_to_csv_test( .is_ok() ); - let written = tmp_store.read(&output_path).await.unwrap(); + let written = tmp_store.read(output_path).await.unwrap(); let format_expect = store.read(format_path).await.unwrap(); assert_eq_lines(written.to_vec(), format_expect.to_vec()); } diff --git a/src/common/datasource/src/tests.rs b/src/common/datasource/src/tests.rs index 8be9add7c88..b983ef86627 100644 --- a/src/common/datasource/src/tests.rs +++ b/src/common/datasource/src/tests.rs @@ -18,34 +18,31 @@ use crate::test_util; #[tokio::test] async fn test_stream_to_json() { - let origin_path = &find_workspace_path("/src/common/datasource/tests/json/basic.json") - .display() - .to_string(); + let root = find_workspace_path("/src/common/datasource/tests/json"); + let store = test_util::test_store(root.to_str().unwrap()); + let origin_path = "basic.json"; // A small threshold // Triggers the flush each writes - test_util::setup_stream_to_json_test(origin_path, |size| size / 2).await; + test_util::setup_stream_to_json_test(&store, origin_path, |size| size / 2).await; // A large threshold // Only triggers the flush at last - test_util::setup_stream_to_json_test(origin_path, |size| size * 2).await; + test_util::setup_stream_to_json_test(&store, origin_path, |size| size * 2).await; } #[tokio::test] async fn test_stream_to_csv() { - let origin_path = &find_workspace_path("/src/common/datasource/tests/csv/basic.csv") - .display() - .to_string(); - - let format_path = &find_workspace_path("/src/common/datasource/tests/csv/basic_format.csv") - .display() - .to_string(); + let root = find_workspace_path("/src/common/datasource/tests/csv"); + let store = test_util::test_store(root.to_str().unwrap()); + let origin_path = "basic.csv"; + let format_path = "basic_format.csv"; // A small threshold // Triggers the flush each writes - test_util::setup_stream_to_csv_test(origin_path, format_path, |size| size / 2).await; + test_util::setup_stream_to_csv_test(&store, origin_path, format_path, |size| size / 2).await; // A large threshold // Only triggers the flush at last - test_util::setup_stream_to_csv_test(origin_path, format_path, |size| size * 2).await; + test_util::setup_stream_to_csv_test(&store, origin_path, format_path, |size| size * 2).await; } diff --git a/src/common/error/src/ext.rs b/src/common/error/src/ext.rs index 59895551156..0c267dc15b2 100644 --- a/src/common/error/src/ext.rs +++ b/src/common/error/src/ext.rs @@ -21,6 +21,20 @@ use std::sync::Arc; use serde::{Deserialize, Deserializer, Serializer}; use snafu::{FromString, Snafu}; +/// Returns the root cause of an error's source chain, i.e. the last error +/// reachable via [`std::error::Error::source`]. For an error without any +/// source the error itself is returned. +/// +/// Mirrors `err.sources().last().unwrap()` (unstable `error_iter`, which +/// yields the error itself followed by its sources). +fn error_chain_root(err: &dyn std::error::Error) -> &dyn std::error::Error { + let mut root = err; + while let Some(source) = root.source() { + root = source; + } + root +} + use crate::status_code::StatusCode; /// Describes whether an error instance is safe and useful to retry. @@ -148,7 +162,7 @@ pub trait ErrorExt: StackError { _ => { let error = self.last(); if let Some(external_error) = error.source() { - let external_root = external_error.sources().last().unwrap(); + let external_root = error_chain_root(external_error); if error.transparent() { format!("{external_root}") @@ -169,7 +183,7 @@ pub trait ErrorExt: StackError { { let error = self.last(); if let Some(external_error) = error.source() { - let external_root = external_error.sources().last().unwrap(); + let external_root = error_chain_root(external_error); Some(external_root) } else { None diff --git a/src/common/error/src/lib.rs b/src/common/error/src/lib.rs index eeedc637133..29ae5c3b1cc 100644 --- a/src/common/error/src/lib.rs +++ b/src/common/error/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(error_iter)] - pub mod ext; pub mod mock; pub mod status_code; diff --git a/src/common/event-recorder/src/lib.rs b/src/common/event-recorder/src/lib.rs index 7ee71722229..0a5d051a01e 100644 --- a/src/common/event-recorder/src/lib.rs +++ b/src/common/event-recorder/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(duration_constructors)] - pub mod context; pub mod error; pub mod event_table; diff --git a/src/common/event-recorder/src/recorder.rs b/src/common/event-recorder/src/recorder.rs index b6c18ed08de..040b2fbcbb3 100644 --- a/src/common/event-recorder/src/recorder.rs +++ b/src/common/event-recorder/src/recorder.rs @@ -108,9 +108,9 @@ fn event_type_filter_is_all(event_types: &EventTypeFilterRef) -> bool { /// The time interval for flushing batched events to the event handler. pub const DEFAULT_FLUSH_INTERVAL_SECONDS: Duration = Duration::from_secs(5); /// The default TTL(90 days) for the events table. -const DEFAULT_EVENTS_TABLE_TTL: Duration = Duration::from_days(90); +const DEFAULT_EVENTS_TABLE_TTL: Duration = Duration::from_secs(90 * 86400); /// The default compaction time window for the events table. -pub const DEFAULT_COMPACTION_TIME_WINDOW: Duration = Duration::from_days(1); +pub const DEFAULT_COMPACTION_TIME_WINDOW: Duration = Duration::from_secs(86400); // The capacity of the tokio channel for transmitting events to background processor. const DEFAULT_CHANNEL_SIZE: usize = 2048; // The size of the buffer for batching events before flushing to event handler. @@ -291,17 +291,29 @@ pub struct EventRecorderOptions { skip_serializing_if = "event_type_filter_is_all" )] pub event_types: EventTypeFilterRef, + /// How long buffered events may wait before reaching the event handler. + /// + /// Not a configuration option: it is skipped by serde so it never appears in + /// config files or `config/config.md`. Integration tests shorten it so they + /// do not pay a flush window per DDL they assert on. + #[serde(skip, default = "default_flush_interval")] + pub flush_interval: Duration, } fn default_events_table_ttl() -> Duration { DEFAULT_EVENTS_TABLE_TTL } +fn default_flush_interval() -> Duration { + DEFAULT_FLUSH_INTERVAL_SECONDS +} + impl Default for EventRecorderOptions { fn default() -> Self { Self { ttl: DEFAULT_EVENTS_TABLE_TTL, event_types: Arc::new(EventTypeFilter::All), + flush_interval: DEFAULT_FLUSH_INTERVAL_SECONDS, } } } @@ -321,13 +333,18 @@ pub struct EventRecorderImpl { impl EventRecorderImpl { pub fn new(event_handler: Box) -> Self { - Self::with_event_type_filter(event_handler, Arc::new(EventTypeFilter::All)) + Self::with_event_type_filter( + event_handler, + Arc::new(EventTypeFilter::All), + DEFAULT_FLUSH_INTERVAL_SECONDS, + ) } /// Creates an event recorder with an event-type filter. pub fn with_event_type_filter( event_handler: Box, event_types: EventTypeFilterRef, + flush_interval: Duration, ) -> Self { let (tx, rx) = channel(DEFAULT_CHANNEL_SIZE); let cancel_token = CancellationToken::new(); @@ -339,13 +356,9 @@ impl EventRecorderImpl { cancel_token: cancel_token.clone(), }; - let processor = EventProcessor::new( - rx, - event_handler, - DEFAULT_FLUSH_INTERVAL_SECONDS, - DEFAULT_MAX_RETRY_TIMES, - ) - .with_cancel_token(cancel_token); + let processor = + EventProcessor::new(rx, event_handler, flush_interval, DEFAULT_MAX_RETRY_TIMES) + .with_cancel_token(cancel_token); // Spawn a background task to process the events. let handle = tokio::spawn(async move { @@ -611,6 +624,7 @@ mod tests { count: count.clone(), }), event_type_filter.clone(), + DEFAULT_FLUSH_INTERVAL_SECONDS, ); assert!(Arc::ptr_eq( diff --git a/src/common/frontend/Cargo.toml b/src/common/frontend/Cargo.toml index 32a5146b566..84a8373dcbb 100644 --- a/src/common/frontend/Cargo.toml +++ b/src/common/frontend/Cargo.toml @@ -12,14 +12,11 @@ common-event-recorder.workspace = true common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true -common-time.workspace = true greptime-proto.workspace = true -humantime.workspace = true +lazy_static.workspace = true meta-client.workspace = true +prometheus.workspace = true serde.workspace = true serde_json.workspace = true snafu.workspace = true tonic.workspace = true - -[dev-dependencies] -tokio.workspace = true diff --git a/src/common/frontend/src/lib.rs b/src/common/frontend/src/lib.rs index dacffd19968..b6ec72da1a3 100644 --- a/src/common/frontend/src/lib.rs +++ b/src/common/frontend/src/lib.rs @@ -18,6 +18,7 @@ use std::str::FromStr; use snafu::OptionExt; pub mod error; +pub mod metrics; pub mod selector; pub mod slow_query_event; diff --git a/src/common/frontend/src/metrics.rs b/src/common/frontend/src/metrics.rs new file mode 100644 index 00000000000..7a2c152fa1f --- /dev/null +++ b/src/common/frontend/src/metrics.rs @@ -0,0 +1,117 @@ +// 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. + +//! Cache lookup counters shared by frontend, pipeline, and protocol handlers. +//! +//! The `type` label identifies `otlp_trace_aux`, `otlp_metrics_legacy`, +//! `pipeline_table`, `pipeline`, `pipeline_original`, `pipeline_failover`, or +//! `mysql_prepared_stmt`. Counters aggregate across catalogs and connections. +//! Trace lookups count distinct service/operation keys once per request, excluding +//! rechecks; legacy metric lookups count each requested metric name. Pipeline +//! loaders count successful coalesced waiters as hits and failed loads as misses. +//! Failover counts one outcome after resolving all candidate schemas; an ambiguous +//! result is a miss. Table and prepared-statement caches count each lookup. +//! Both outcomes for all known caches are exported from the first metric update. + +use lazy_static::lazy_static; +use prometheus::{IntCounterVec, register_int_counter_vec}; + +const CACHE_TYPES: [&str; 7] = [ + "otlp_trace_aux", + "otlp_metrics_legacy", + "pipeline_table", + "pipeline", + "pipeline_original", + "pipeline_failover", + "mysql_prepared_stmt", +]; + +lazy_static! { + static ref CACHE_COUNTERS: (IntCounterVec, IntCounterVec) = { + let hits = register_int_counter_vec!( + "greptime_frontend_cache_hit", + "frontend cache hit", + &["type"] + ) + .unwrap(); + let misses = register_int_counter_vec!( + "greptime_frontend_cache_miss", + "frontend cache miss", + &["type"] + ) + .unwrap(); + // Register both outcomes together so an all-hit or all-miss cache has both series. + for cache_type in CACHE_TYPES { + hits.with_label_values(&[cache_type]); + misses.with_label_values(&[cache_type]); + } + (hits, misses) + }; + + /// Frontend cache hits by cache type. + pub static ref CACHE_HIT: IntCounterVec = CACHE_COUNTERS.0.clone(); + /// Frontend cache misses by cache type. + pub static ref CACHE_MISS: IntCounterVec = CACHE_COUNTERS.1.clone(); +} + +/// Records one completed cache lookup. +pub fn record_cache_lookup(cache_type: &str, hit: bool) { + if hit { + CACHE_HIT.with_label_values(&[cache_type]).inc(); + } else { + CACHE_MISS.with_label_values(&[cache_type]).inc(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cache_lookup_counters() { + record_cache_lookup("pipeline", false); + // Scrape before accessing CACHE_HIT, which would initialize it and hide a missing series. + let families = prometheus::gather(); + for (name, expected) in [ + ("greptime_frontend_cache_hit", 0.0), + ("greptime_frontend_cache_miss", 1.0), + ] { + let family = families + .iter() + .find(|family| family.name() == name) + .unwrap(); + assert_eq!(family.get_metric().len(), CACHE_TYPES.len()); + for cache_type in CACHE_TYPES { + let metric = family + .get_metric() + .iter() + .find(|metric| metric.get_label()[0].value() == cache_type) + .unwrap(); + assert_eq!(metric.get_label()[0].name(), "type"); + assert_eq!( + metric.get_counter().value(), + if cache_type == "pipeline" { + expected + } else { + 0.0 + } + ); + } + } + record_cache_lookup("pipeline", true); + record_cache_lookup("pipeline", false); + assert_eq!(CACHE_HIT.with_label_values(&["pipeline"]).get(), 1); + assert_eq!(CACHE_MISS.with_label_values(&["pipeline"]).get(), 2); + } +} diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index ee6cfac9dfb..3a5da30c136 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -17,21 +17,17 @@ ai_functions = ["dep:futures", "dep:reqwest"] [dependencies] ahash.workspace = true api.workspace = true -arc-swap = "1.0" arrow.workspace = true -arrow-cast.workspace = true arrow-schema.workspace = true async-trait.workspace = true bincode = "=1.3.3" catalog.workspace = true -chrono.workspace = true common-base.workspace = true common-catalog.workspace = true common-error.workspace = true common-macro.workspace = true common-meta.workspace = true common-query.workspace = true -common-runtime.workspace = true common-telemetry.workspace = true common-time.workspace = true common-version.workspace = true @@ -55,7 +51,6 @@ jsonb.workspace = true memchr = "2.7" mito-codec.workspace = true nalgebra.workspace = true -num = "0.4" num-traits = "0.2" paste.workspace = true regex.workspace = true diff --git a/src/common/function/src/admin.rs b/src/common/function/src/admin.rs index ff8172ebb5a..79182dd74e9 100644 --- a/src/common/function/src/admin.rs +++ b/src/common/function/src/admin.rs @@ -34,7 +34,7 @@ use reconcile_catalog::ReconcileCatalogFunction; use reconcile_database::ReconcileDatabaseFunction; use reconcile_table::ReconcileTableFunction; -use crate::admin::build_index_table::BuildIndexFunction; +use crate::admin::build_index_table::{BuildIndexFunction, BuildSeriesIndexFunction}; use crate::admin::discard_unflushed_data::DiscardUnflushedDataFunction; use crate::flush_flow::FlushFlowFunction; use crate::function_registry::FunctionRegistry; @@ -53,6 +53,7 @@ impl AdminFunction { registry.register(GcRegionsFunction::factory()); registry.register(GcTableFunction::factory()); registry.register(BuildIndexFunction::factory()); + registry.register(BuildSeriesIndexFunction::factory()); registry.register(FlushFlowFunction::factory()); registry.register(ReconcileCatalogFunction::factory()); registry.register(ReconcileDatabaseFunction::factory()); diff --git a/src/common/function/src/admin/build_index_table.rs b/src/common/function/src/admin/build_index_table.rs index 155f198c792..24026aead8d 100644 --- a/src/common/function/src/admin/build_index_table.rs +++ b/src/common/function/src/admin/build_index_table.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use api::v1::region::build_index_request; use arrow::datatypes::DataType as ArrowDataType; use common_error::ext::BoxedError; use common_macro::admin_fn; @@ -38,6 +39,45 @@ pub(crate) async fn build_index( table_mutation_handler: &TableMutationHandlerRef, query_ctx: &QueryContextRef, params: &[ValueRef<'_>], +) -> Result { + build_index_impl( + table_mutation_handler, + query_ctx, + params, + "build_index", + None, + ) + .await +} + +/// Reconciles series indexes on every physical data-region leader and waits for publication. +#[admin_fn( + name = BuildSeriesIndexFunction, + display_name = build_series_index, + sig_fn = build_index_signature, + ret = uint64 +)] +pub(crate) async fn build_series_index( + table_mutation_handler: &TableMutationHandlerRef, + query_ctx: &QueryContextRef, + params: &[ValueRef<'_>], +) -> Result { + build_index_impl( + table_mutation_handler, + query_ctx, + params, + "build_series_index", + Some(build_index_request::Options::SeriesIndex(Default::default())), + ) + .await +} + +async fn build_index_impl( + table_mutation_handler: &TableMutationHandlerRef, + query_ctx: &QueryContextRef, + params: &[ValueRef<'_>], + function: &str, + options: Option, ) -> Result { ensure!( params.len() == 1, @@ -51,7 +91,7 @@ pub(crate) async fn build_index( let ValueRef::String(table_name) = params[0] else { return UnsupportedInputDataTypeSnafu { - function: "build_index", + function, datatypes: params.iter().map(|v| v.data_type()).collect::>(), } .fail(); @@ -64,6 +104,7 @@ pub(crate) async fn build_index( let affected_rows = table_mutation_handler .build_index( BuildIndexTableRequest { + options, catalog_name, schema_name, table_name, @@ -78,3 +119,45 @@ pub(crate) async fn build_index( fn build_index_signature() -> Signature { Signature::uniform(1, vec![ArrowDataType::Utf8], Volatility::Immutable) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::function::FunctionContext; + + #[tokio::test] + async fn test_build_series_index_arguments() { + let ctx = FunctionContext::mock(); + let handler = ctx.state.table_mutation_handler.as_ref().unwrap(); + for params in [ + vec![], + vec![ValueRef::Int32(1)], + vec![ValueRef::Null], + vec![ValueRef::String("a"), ValueRef::String("b")], + ] { + assert!( + build_index_impl( + handler, + &ctx.query_ctx, + ¶ms, + "build_series_index", + Some(build_index_request::Options::SeriesIndex(Default::default())) + ) + .await + .is_err() + ); + } + assert_eq!( + Value::UInt64(42), + build_index_impl( + handler, + &ctx.query_ctx, + &[ValueRef::String("greptime.public.test")], + "build_series_index", + Some(build_index_request::Options::SeriesIndex(Default::default())), + ) + .await + .unwrap() + ); + } +} diff --git a/src/common/function/src/aggrs/aggr_wrapper.rs b/src/common/function/src/aggrs/aggr_wrapper.rs index 70651eee41d..710c383a834 100644 --- a/src/common/function/src/aggrs/aggr_wrapper.rs +++ b/src/common/function/src/aggrs/aggr_wrapper.rs @@ -25,7 +25,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, StructArray}; +use arrow::array::{ArrayData, ArrayRef, BooleanArray, StructArray, make_array}; use arrow_schema::{FieldRef, Fields}; use common_telemetry::debug; use datafusion::functions_aggregate::all_default_aggregate_functions; @@ -84,6 +84,19 @@ pub fn is_all_aggr_exprs_steppable(aggr_exprs: &[Expr]) -> bool { return false; } + // DataFusion only sorts the input of an aggregate with a hard ordering requirement + // when the requirement is already satisfied or the aggregate has a reverse + // expression (apache/datafusion#25676). The state wrapper has none, so e.g. + // `nth_value(.. ORDER BY ..)` would read unsorted input on datanodes. Ordered-set + // aggregates like `approx_percentile_cont(..) WITHIN GROUP (ORDER BY ..)` are + // exempt: their ORDER BY names the value, and they don't need sorted input. + if !aggr_func.params.order_by.is_empty() + && aggr_func.func.order_sensitivity().hard_requires() + && !aggr_func.func.supports_within_group_clause() + { + return false; + } + // whether the corresponding state function exists in the registry FUNCTION_REGISTRY.is_aggr_func_exist(&aggr_state_func_name(aggr_func.func.name())) } else { @@ -526,43 +539,7 @@ impl StateGroupsAccum { } fn wrap_state_arrays(&self, arrays: Vec) -> datafusion_common::Result { - let array_type = arrays - .iter() - .map(|array| array.data_type().clone()) - .collect::>(); - let expected_type = self - .state_fields - .iter() - .map(|field| field.data_type().clone()) - .collect::>(); - if array_type != expected_type { - debug!( - "State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}", - self.state_fields.len(), - arrays.len(), - self.state_fields, - array_type, - ); - let guess_schema = arrays - .iter() - .enumerate() - .map(|(index, array)| { - Field::new( - format!("col_{index}[mismatch_state]").as_str(), - array.data_type().clone(), - true, - ) - }) - .collect::(); - let array = StructArray::try_new(guess_schema, arrays, None)?; - return Ok(Arc::new(array)); - } - - Ok(Arc::new(StructArray::try_new( - self.state_fields.clone(), - arrays, - None, - )?)) + Ok(Arc::new(state_struct_array(&self.state_fields, arrays)?)) } } @@ -610,6 +587,81 @@ impl GroupsAccumulator for StateGroupsAccum { } } +/// Wraps the state arrays of an accumulator into a struct of the declared state fields. +/// +/// The declared state type is derived from logical expressions, while the accumulator names +/// nested fields after physical expressions. For example, `array_agg(v ORDER BY ts)` declares +/// its orderings as `List(Struct("ts": ..))` but produces `List(Struct("ts@0": ..))`. Arrays that +/// differ only in nested field names are cast to the declared type; any other difference is an +/// error. +fn state_struct_array( + state_fields: &Fields, + arrays: Vec, +) -> datafusion_common::Result { + if arrays.len() != state_fields.len() { + return Err(datafusion_common::DataFusionError::Internal(format!( + "Expected {} state arrays for fields {:?}, got {}", + state_fields.len(), + state_fields, + arrays.len() + ))); + } + let arrays = arrays + .into_iter() + .zip(state_fields.iter()) + .map(|(array, field)| { + let expected = field.data_type(); + if array.data_type() == expected { + Ok(array) + } else if array.data_type().equals_datatype(expected) { + Ok(make_array(relabel_nested_fields( + array.to_data(), + expected, + )?)) + } else { + Err(datafusion_common::DataFusionError::Internal(format!( + "State field `{}` expects type {expected}, but the accumulator produced {}", + field.name(), + array.data_type() + ))) + } + }) + .collect::>>()?; + Ok(StructArray::try_new(state_fields.clone(), arrays, None)?) +} + +/// Rebuilds `data` with the type `target`, which must match it position by position apart +/// from nested field names and metadata. Unlike a cast, children are never matched by name. +fn relabel_nested_fields( + data: ArrayData, + target: &DataType, +) -> datafusion_common::Result { + let child_types = match target { + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => vec![field.data_type()], + DataType::Struct(fields) => fields.iter().map(|f| f.data_type()).collect(), + _ if data.child_data().is_empty() => vec![], + _ => { + return Err(datafusion_common::DataFusionError::NotImplemented(format!( + "Relabeling nested fields of {target}" + ))); + } + }; + let children = data + .child_data() + .iter() + .zip(child_types) + .map(|(child, child_type)| relabel_nested_fields(child.clone(), child_type)) + .collect::>>()?; + Ok(data + .into_builder() + .data_type(target.clone()) + .child_data(children) + .build()?) +} + impl StateAccum { pub fn new( inner: Box, @@ -636,40 +688,7 @@ impl Accumulator for StateAccum { .iter() .map(|s| s.to_array()) .collect::, _>>()?; - let array_type = array - .iter() - .map(|a| a.data_type().clone()) - .collect::>(); - let expected_type: Vec<_> = self - .state_fields - .iter() - .map(|f| f.data_type().clone()) - .collect(); - if array_type != expected_type { - debug!( - "State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}", - self.state_fields.len(), - array.len(), - self.state_fields, - array_type, - ); - let guess_schema = array - .iter() - .enumerate() - .map(|(index, array)| { - Field::new( - format!("col_{index}[mismatch_state]").as_str(), - array.data_type().clone(), - true, - ) - }) - .collect::(); - let arr = StructArray::try_new(guess_schema, array, None)?; - - return Ok(ScalarValue::Struct(Arc::new(arr))); - } - - let struct_array = StructArray::try_new(self.state_fields.clone(), array, None)?; + let struct_array = state_struct_array(&self.state_fields, array)?; Ok(ScalarValue::Struct(Arc::new(struct_array))) } diff --git a/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs b/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs index 480b0b5a490..715b051407d 100644 --- a/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs +++ b/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs @@ -153,13 +153,17 @@ fn rewrite_expr( if is_fix { // then always fix the ordering field&distinct flag and more let order_by = aggregate_function.params.order_by.clone(); + // DataFusion always makes ordering fields nullable when building the physical + // aggregate (`ordering_fields` in `datafusion-functions-aggregate-common`), and some + // accumulators (e.g. `array_agg`) nest these fields in their state type. Keep the + // declared state type in line with what the accumulator produces. let ordering_fields: Vec<_> = order_by .iter() .map(|sort_expr| { sort_expr .expr .to_field(&aggregate_input.schema()) - .map(|(_, f)| f) + .map(|(_, f)| Arc::new(f.as_ref().clone().with_nullable(true))) }) .collect::>>()?; let distinct = aggregate_function.params.distinct; diff --git a/src/common/function/src/aggrs/aggr_wrapper/tests.rs b/src/common/function/src/aggrs/aggr_wrapper/tests.rs index 4274a03cd89..b68a735f182 100644 --- a/src/common/function/src/aggrs/aggr_wrapper/tests.rs +++ b/src/common/function/src/aggrs/aggr_wrapper/tests.rs @@ -27,8 +27,11 @@ use datafusion::catalog::{Session, TableProvider}; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::datasource::DefaultTableSource; use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion::functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; +use datafusion::functions_aggregate::array_agg::array_agg_udaf; use datafusion::functions_aggregate::average::avg_udaf; use datafusion::functions_aggregate::count::count_udaf; +use datafusion::functions_aggregate::nth_value::nth_value_udaf; use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::optimizer::AnalyzerRule; use datafusion::optimizer::analyzer::type_coercion::TypeCoercion; @@ -1289,6 +1292,40 @@ async fn test_udaf_correct_eval_result() { order_by: vec![], null_treatment: None, }, + // The ordering state of `array_agg` nests the ORDER BY fields in `List(Struct(..))`. + TestCase { + func: array_agg_udaf(), + input_schema: Arc::new(arrow_schema::Schema::new(vec![ + Field::new("number", DataType::Float64, true), + Field::new( + "ts", + DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + false, + ), + ])), + args: vec![Expr::Column(Column::new_unqualified("number"))], + input: vec![ + Arc::new(Float64Array::from(vec![Some(3.), Some(1.), Some(2.)])), + Arc::new(TimestampMillisecondArray::from(vec![3000, 1000, 2000])), + ], + expected_output: Some(ScalarValue::List(ScalarValue::new_list_nullable( + &[ + ScalarValue::Float64(Some(1.)), + ScalarValue::Float64(Some(2.)), + ScalarValue::Float64(Some(3.)), + ], + &DataType::Float64, + ))), + expected_fn: None, + distinct: false, + filter: None, + order_by: vec![SortExpr::new( + Expr::Column(Column::new_unqualified("ts")), + true, + true, + )], + null_treatment: None, + }, // TODO(discord9): udd_merge/hll_merge/geo_path/quantile_aggr tests ]; let test_table_ref = TableReference::bare("TestTable"); @@ -1396,3 +1433,89 @@ async fn execute_phy_plan( } Ok(batches) } + +#[test] +fn test_state_struct_array_rejects_mismatched_state() { + let fields = Fields::from(vec![Field::new("sum", DataType::Int64, true)]); + let arrays: Vec = vec![Arc::new(Float64Array::from(vec![1.0]))]; + let err = state_struct_array(&fields, arrays).unwrap_err(); + assert!( + err.to_string() + .contains("State field `sum` expects type Int64, but the accumulator produced Float64"), + "{err}" + ); +} + +#[test] +fn test_state_struct_array_keeps_child_order() { + let int_field = |name: &str| Field::new(name, DataType::Int64, true); + // Same child types, names crossed: children must stay in place, not be matched by name. + let produced = StructArray::from(vec![ + ( + Arc::new(int_field("a")), + Arc::new(Int64Array::from(vec![10])) as ArrayRef, + ), + ( + Arc::new(int_field("b")), + Arc::new(Int64Array::from(vec![20])) as ArrayRef, + ), + ]); + let declared = Fields::from(vec![Field::new( + "state", + DataType::Struct(Fields::from(vec![int_field("b"), int_field("a")])), + true, + )]); + + let state = state_struct_array(&declared, vec![Arc::new(produced)]).unwrap(); + let state = state.column(0).as_struct(); + assert_eq!( + state + .column_by_name("b") + .unwrap() + .as_primitive::() + .value(0), + 10 + ); + assert_eq!( + state + .column_by_name("a") + .unwrap() + .as_primitive::() + .value(0), + 20 + ); +} + +#[test] +fn test_hard_ordered_aggr_not_steppable() { + let order_by = vec![SortExpr::new( + Expr::Column(Column::new_unqualified("ts")), + true, + true, + )]; + let aggr = |func: Arc, args: Vec| { + Expr::AggregateFunction(AggregateFunction::new_udf( + func, + args, + false, + None, + order_by.clone(), + None, + )) + }; + let number = Expr::Column(Column::new_unqualified("number")); + + assert!(!is_all_aggr_exprs_steppable(&[aggr( + nth_value_udaf(), + vec![number.clone(), lit(2i64)], + )])); + assert!(is_all_aggr_exprs_steppable(&[aggr( + array_agg_udaf(), + vec![number.clone()] + )])); + // WITHIN GROUP (ORDER BY number) + assert!(is_all_aggr_exprs_steppable(&[aggr( + approx_percentile_cont_udaf(), + vec![number, lit(0.5f64)] + )])); +} diff --git a/src/common/function/src/aggrs/geo.rs b/src/common/function/src/aggrs/geo.rs index c76fe7348f2..abe9070a0a0 100644 --- a/src/common/function/src/aggrs/geo.rs +++ b/src/common/function/src/aggrs/geo.rs @@ -21,7 +21,7 @@ pub(crate) struct GeoFunction; impl GeoFunction { pub fn register(registry: &FunctionRegistry) { - registry.register_aggr(encoding::JsonEncodePathAccumulator::uadf_impl()); - registry.register_aggr(geo_path::GeoPathAccumulator::uadf_impl()); + registry.register_aggr(encoding::JsonEncodePathAccumulator::udf_impl()); + registry.register_aggr(geo_path::GeoPathAccumulator::udf_impl()); } } diff --git a/src/common/function/src/aggrs/geo/encoding.rs b/src/common/function/src/aggrs/geo/encoding.rs index 35a59d1ecfb..5b4591bb6d8 100644 --- a/src/common/function/src/aggrs/geo/encoding.rs +++ b/src/common/function/src/aggrs/geo/encoding.rs @@ -47,7 +47,7 @@ impl JsonEncodePathAccumulator { Self::default() } - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { create_udaf( JSON_ENCODE_PATH_NAME, // Input types: lat, lng, timestamp @@ -221,18 +221,24 @@ impl DfAccumulator for JsonEncodePathAccumulator { ))); } - for state in states { - let state = as_struct_array(state)?; - let lat_list = as_list_array(state.column(0))?.value(0); - let lat_array = as_primitive_array::(&lat_list)?; - let lng_list = as_list_array(state.column(1))?.value(0); - let lng_array = as_primitive_array::(&lng_list)?; - let ts_list = as_list_array(state.column(2))?.value(0); - let ts_array = as_primitive_array::(&ts_list)?; + let state = as_struct_array(&states[0])?; + let lat_lists = as_list_array(state.column(0))?; + let lng_lists = as_list_array(state.column(1))?; + let ts_lists = as_list_array(state.column(2))?; + for row in 0..state.len() { + if state.is_null(row) { + continue; + } + let lat_list = lat_lists.value(row); + let lng_list = lng_lists.value(row); + let ts_list = ts_lists.value(row); - self.lat.extend(lat_array); - self.lng.extend(lng_array); - self.timestamp.extend(ts_array); + self.lat + .extend(as_primitive_array::(&lat_list)?); + self.lng + .extend(as_primitive_array::(&lng_list)?); + self.timestamp + .extend(as_primitive_array::(&ts_list)?); } Ok(()) @@ -331,9 +337,9 @@ mod tests { _ => panic!("Expected Struct scalar value"), }; - // Merge state arrays - merged.merge_batch(&[state_array1]).unwrap(); - merged.merge_batch(&[state_array2]).unwrap(); + // States from different partitions arrive as rows of one batch + let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap(); + merged.merge_batch(&[states]).unwrap(); // Evaluate merged result let result = merged.evaluate().unwrap(); diff --git a/src/common/function/src/aggrs/geo/geo_path.rs b/src/common/function/src/aggrs/geo/geo_path.rs index b09745eaeb7..fe8d8bae21c 100644 --- a/src/common/function/src/aggrs/geo/geo_path.rs +++ b/src/common/function/src/aggrs/geo/geo_path.rs @@ -47,7 +47,7 @@ impl GeoPathAccumulator { Self::default() } - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { create_udaf( GEO_PATH_NAME, // Input types: lat, lng, timestamp @@ -250,18 +250,24 @@ impl DfAccumulator for GeoPathAccumulator { ))); } - for state in states { - let state = as_struct_array(state)?; - let lat_list = as_list_array(state.column(0))?.value(0); - let lat_array = as_primitive_array::(&lat_list)?; - let lng_list = as_list_array(state.column(1))?.value(0); - let lng_array = as_primitive_array::(&lng_list)?; - let ts_list = as_list_array(state.column(2))?.value(0); - let ts_array = as_primitive_array::(&ts_list)?; + let state = as_struct_array(&states[0])?; + let lat_lists = as_list_array(state.column(0))?; + let lng_lists = as_list_array(state.column(1))?; + let ts_lists = as_list_array(state.column(2))?; + for row in 0..state.len() { + if state.is_null(row) { + continue; + } + let lat_list = lat_lists.value(row); + let lng_list = lng_lists.value(row); + let ts_list = ts_lists.value(row); - self.lat.extend(lat_array); - self.lng.extend(lng_array); - self.timestamp.extend(ts_array); + self.lat + .extend(as_primitive_array::(&lat_list)?); + self.lng + .extend(as_primitive_array::(&lng_list)?); + self.timestamp + .extend(as_primitive_array::(&ts_list)?); } Ok(()) @@ -403,9 +409,9 @@ mod tests { _ => panic!("Expected Struct scalar value"), }; - // Merge state arrays - merged.merge_batch(&[state_array1]).unwrap(); - merged.merge_batch(&[state_array2]).unwrap(); + // States from different partitions arrive as rows of one batch + let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap(); + merged.merge_batch(&[states]).unwrap(); // Evaluate merged result let result = merged.evaluate().unwrap(); diff --git a/src/common/function/src/aggrs/vector.rs b/src/common/function/src/aggrs/vector.rs index 03489a51d44..0c0c70f7360 100644 --- a/src/common/function/src/aggrs/vector.rs +++ b/src/common/function/src/aggrs/vector.rs @@ -25,8 +25,8 @@ pub(crate) struct VectorFunction; impl VectorFunction { pub fn register(registry: &FunctionRegistry) { - registry.register_aggr(VectorSum::uadf_impl()); - registry.register_aggr(VectorProduct::uadf_impl()); - registry.register_aggr(VectorAvg::uadf_impl()); + registry.register_aggr(VectorSum::udf_impl()); + registry.register_aggr(VectorProduct::udf_impl()); + registry.register_aggr(VectorAvg::udf_impl()); } } diff --git a/src/common/function/src/aggrs/vector/avg.rs b/src/common/function/src/aggrs/vector/avg.rs index 822a7c16efe..a8af8cf6d0c 100644 --- a/src/common/function/src/aggrs/vector/avg.rs +++ b/src/common/function/src/aggrs/vector/avg.rs @@ -39,7 +39,7 @@ pub struct VectorAvg { impl VectorAvg { /// Create a new `AggregateUDF` for the `vec_avg` aggregate function. - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { let signature = Signature::one_of( vec![ TypeSignature::Exact(vec![DataType::Utf8]), @@ -62,14 +62,14 @@ impl VectorAvg { } fn accumulator(args: AccumulatorArgs) -> Result> { - if args.schema.fields().len() != 1 { + if args.exprs.len() != 1 { return Err(datafusion_common::DataFusionError::Internal(format!( "expect creating `VEC_AVG` with only one input field, actual {}", - args.schema.fields().len() + args.exprs.len() ))); } - let t = args.schema.field(0).data_type(); + let t = args.expr_fields[0].data_type(); if !matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary) { return Err(datafusion_common::DataFusionError::Internal(format!( "unexpected input datatype {t} when creating `VEC_AVG`" diff --git a/src/common/function/src/aggrs/vector/product.rs b/src/common/function/src/aggrs/vector/product.rs index 11cfa4cb1a8..4f0e9bb6217 100644 --- a/src/common/function/src/aggrs/vector/product.rs +++ b/src/common/function/src/aggrs/vector/product.rs @@ -28,6 +28,9 @@ use crate::scalars::vector::impl_conv::{ }; /// Aggregates by multiplying elements across the same dimension, returns a vector. +/// +/// The result is NULL if any input vector is NULL, so the partial state carries a +/// `has_null` flag: a NULL `product` alone can't tell a NULL input from an empty partition. #[derive(Debug, Default)] pub struct VectorProduct { product: Option>, @@ -36,7 +39,7 @@ pub struct VectorProduct { impl VectorProduct { /// Create a new `AggregateUDF` for the `vec_product` aggregate function. - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { let signature = Signature::one_of( vec![ TypeSignature::Exact(vec![DataType::Utf8]), @@ -49,20 +52,23 @@ impl VectorProduct { signature, DataType::Binary, Arc::new(Self::accumulator), - vec![Arc::new(Field::new("x", DataType::Binary, true))], + vec![ + Arc::new(Field::new("product", DataType::Binary, true)), + Arc::new(Field::new("has_null", DataType::Boolean, true)), + ], ); AggregateUDF::from(udaf) } fn accumulator(args: AccumulatorArgs) -> Result> { - if args.schema.fields().len() != 1 { + if args.exprs.len() != 1 { return Err(datafusion_common::DataFusionError::Internal(format!( "expect creating `VEC_PRODUCT` with only one input field, actual {}", - args.schema.fields().len() + args.exprs.len() ))); } - let t = args.schema.field(0).data_type(); + let t = args.expr_fields[0].data_type(); if !matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary) { return Err(datafusion_common::DataFusionError::Internal(format!( "unexpected input datatype {t} when creating `VEC_PRODUCT`" @@ -72,13 +78,33 @@ impl VectorProduct { Ok(Box::new(VectorProduct::default())) } - fn inner(&mut self, len: usize) -> &mut OVector { - self.product.get_or_insert_with(|| { - OVector::from_iterator_generic(Dyn(len), Const::<1>, (0..len).map(|_| 1.0)) - }) + fn mul(&mut self, vector: &[f32]) { + let vector = DVectorView::from_slice(vector, vector.len()); + let product = self.product.get_or_insert_with(|| { + OVector::from_iterator_generic( + Dyn(vector.len()), + Const::<1>, + (0..vector.len()).map(|_| 1.0), + ) + }); + *product = product.component_mul(&vector); } - fn update(&mut self, values: &[ArrayRef], is_update: bool) -> Result<()> { + fn set_null(&mut self) { + self.has_null = true; + self.product = None; + } +} + +impl Accumulator for VectorProduct { + fn state(&mut self) -> Result> { + Ok(vec![ + self.evaluate()?, + ScalarValue::Boolean(Some(self.has_null)), + ]) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { if values.is_empty() || self.has_null { return Ok(()); }; @@ -112,33 +138,34 @@ impl VectorProduct { } }; if vectors.len() != values[0].len() { - if is_update { - self.has_null = true; - self.product = None; - } + self.set_null(); return Ok(()); } - vectors.iter().for_each(|v| { - let v = DVectorView::from_slice(v, v.len()); - let inner = self.inner(v.len()); - *inner = inner.component_mul(&v); - }); + vectors.iter().for_each(|v| self.mul(v)); Ok(()) } -} - -impl Accumulator for VectorProduct { - fn state(&mut self) -> Result> { - self.evaluate().map(|v| vec![v]) - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.update(values, true) - } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - self.update(states, false) + let [products, has_nulls] = states else { + return Err(datafusion_common::DataFusionError::Internal(format!( + "expect 2 states for `VEC_PRODUCT`, actual {}", + states.len() + ))); + }; + if self.has_null { + return Ok(()); + } + if has_nulls.as_boolean().true_count() > 0 { + self.set_null(); + return Ok(()); + } + + // A NULL product without `has_null` comes from a partition without input rows. + for b in products.as_binary::().iter().flatten() { + self.mul(&binlit_as_veclit(b)?); + } + Ok(()) } fn evaluate(&mut self) -> Result { @@ -225,4 +252,42 @@ mod tests { vec_product.evaluate().unwrap() ); } + + #[test] + fn test_merge_batch() { + let partial = |v: Option<&str>| { + let mut acc = VectorProduct::default(); + let v: ArrayRef = Arc::new(StringArray::from(vec![v])); + acc.update_batch(&[v]).unwrap(); + acc.state().unwrap() + }; + let states = |states: Vec>| -> Vec { + (0..2) + .map(|i| ScalarValue::iter_to_array(states.iter().map(|s| s[i].clone())).unwrap()) + .collect() + }; + + // An empty partition in the middle of the batch must not stop the merge. + let mut merged = VectorProduct::default(); + merged + .merge_batch(&states(vec![ + partial(Some("[1.0,2.0]")), + VectorProduct::default().state().unwrap(), + partial(Some("[3.0,4.0]")), + ])) + .unwrap(); + assert_eq!( + ScalarValue::Binary(Some(veclit_to_binlit(&[3.0, 8.0]))), + merged.evaluate().unwrap() + ); + + // A NULL input in any partition makes the result NULL. + merged + .merge_batch(&states(vec![partial(Some("[1.0,2.0]")), partial(None)])) + .unwrap(); + merged + .merge_batch(&states(vec![partial(Some("[3.0,4.0]"))])) + .unwrap(); + assert_eq!(ScalarValue::Binary(None), merged.evaluate().unwrap()); + } } diff --git a/src/common/function/src/aggrs/vector/sum.rs b/src/common/function/src/aggrs/vector/sum.rs index 0ee8f0df7bf..466fc42cccb 100644 --- a/src/common/function/src/aggrs/vector/sum.rs +++ b/src/common/function/src/aggrs/vector/sum.rs @@ -28,6 +28,9 @@ use crate::scalars::vector::impl_conv::{ }; /// The accumulator for the `vec_sum` aggregate function. +/// +/// The result is NULL if any input vector is NULL, so the partial state carries a +/// `has_null` flag: a NULL `sum` alone can't tell a NULL input from an empty partition. #[derive(Debug, Default)] pub struct VectorSum { sum: Option>, @@ -36,7 +39,7 @@ pub struct VectorSum { impl VectorSum { /// Create a new `AggregateUDF` for the `vec_sum` aggregate function. - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { let signature = Signature::one_of( vec![ TypeSignature::Exact(vec![DataType::Utf8]), @@ -49,20 +52,23 @@ impl VectorSum { signature, DataType::Binary, Arc::new(Self::accumulator), - vec![Arc::new(Field::new("x", DataType::Binary, true))], + vec![ + Arc::new(Field::new("sum", DataType::Binary, true)), + Arc::new(Field::new("has_null", DataType::Boolean, true)), + ], ); AggregateUDF::from(udaf) } fn accumulator(args: AccumulatorArgs) -> Result> { - if args.schema.fields().len() != 1 { + if args.exprs.len() != 1 { return Err(datafusion_common::DataFusionError::Internal(format!( "expect creating `VEC_SUM` with only one input field, actual {}", - args.schema.fields().len() + args.exprs.len() ))); } - let t = args.schema.field(0).data_type(); + let t = args.expr_fields[0].data_type(); if !matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary) { return Err(datafusion_common::DataFusionError::Internal(format!( "unexpected input datatype {t} when creating `VEC_SUM`" @@ -72,12 +78,28 @@ impl VectorSum { Ok(Box::new(VectorSum::default())) } - fn inner(&mut self, len: usize) -> &mut OVector { - self.sum - .get_or_insert_with(|| OVector::zeros_generic(Dyn(len), Const::<1>)) + fn add(&mut self, vector: &[f32]) { + let vector = DVectorView::from_slice(vector, vector.len()); + *self + .sum + .get_or_insert_with(|| OVector::zeros_generic(Dyn(vector.len()), Const::<1>)) += vector; } - fn update(&mut self, values: &[ArrayRef], is_update: bool) -> Result<()> { + fn set_null(&mut self) { + self.has_null = true; + self.sum = None; + } +} + +impl Accumulator for VectorSum { + fn state(&mut self) -> Result> { + Ok(vec![ + self.evaluate()?, + ScalarValue::Boolean(Some(self.has_null)), + ]) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { if values.is_empty() || self.has_null { return Ok(()); }; @@ -87,45 +109,30 @@ impl VectorSum { let arr: &StringArray = values[0].as_string(); for s in arr.iter() { let Some(s) = s else { - if is_update { - self.has_null = true; - self.sum = None; - } + self.set_null(); return Ok(()); }; - let values = parse_veclit_from_strlit(s)?; - let vec_column = DVectorView::from_slice(&values, values.len()); - *self.inner(vec_column.len()) += vec_column; + self.add(&parse_veclit_from_strlit(s)?); } } DataType::LargeUtf8 => { let arr: &LargeStringArray = values[0].as_string(); for s in arr.iter() { let Some(s) = s else { - if is_update { - self.has_null = true; - self.sum = None; - } + self.set_null(); return Ok(()); }; - let values = parse_veclit_from_strlit(s)?; - let vec_column = DVectorView::from_slice(&values, values.len()); - *self.inner(vec_column.len()) += vec_column; + self.add(&parse_veclit_from_strlit(s)?); } } DataType::Binary => { let arr: &BinaryArray = values[0].as_binary(); for b in arr.iter() { let Some(b) = b else { - if is_update { - self.has_null = true; - self.sum = None; - } + self.set_null(); return Ok(()); }; - let values = binlit_as_veclit(b)?; - let vec_column = DVectorView::from_slice(&values, values.len()); - *self.inner(vec_column.len()) += vec_column; + self.add(&binlit_as_veclit(b)?); } } _ => { @@ -137,19 +144,27 @@ impl VectorSum { } Ok(()) } -} - -impl Accumulator for VectorSum { - fn state(&mut self) -> Result> { - self.evaluate().map(|v| vec![v]) - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.update(values, true) - } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - self.update(states, false) + let [sums, has_nulls] = states else { + return Err(datafusion_common::DataFusionError::Internal(format!( + "expect 2 states for `VEC_SUM`, actual {}", + states.len() + ))); + }; + if self.has_null { + return Ok(()); + } + if has_nulls.as_boolean().true_count() > 0 { + self.set_null(); + return Ok(()); + } + + // A NULL sum without `has_null` comes from a partition without input rows. + for b in sums.as_binary::().iter().flatten() { + self.add(&binlit_as_veclit(b)?); + } + Ok(()) } fn evaluate(&mut self) -> Result { @@ -236,4 +251,42 @@ mod tests { vec_sum.evaluate().unwrap() ); } + + #[test] + fn test_merge_batch() { + let partial = |v: Option<&str>| { + let mut acc = VectorSum::default(); + let v: ArrayRef = Arc::new(StringArray::from(vec![v])); + acc.update_batch(&[v]).unwrap(); + acc.state().unwrap() + }; + let states = |states: Vec>| -> Vec { + (0..2) + .map(|i| ScalarValue::iter_to_array(states.iter().map(|s| s[i].clone())).unwrap()) + .collect() + }; + + // An empty partition in the middle of the batch must not stop the merge. + let mut merged = VectorSum::default(); + merged + .merge_batch(&states(vec![ + partial(Some("[1.0,2.0]")), + VectorSum::default().state().unwrap(), + partial(Some("[3.0,4.0]")), + ])) + .unwrap(); + assert_eq!( + ScalarValue::Binary(Some(veclit_to_binlit(&[4.0, 6.0]))), + merged.evaluate().unwrap() + ); + + // A NULL input in any partition makes the result NULL. + merged + .merge_batch(&states(vec![partial(Some("[1.0,2.0]")), partial(None)])) + .unwrap(); + merged + .merge_batch(&states(vec![partial(Some("[3.0,4.0]"))])) + .unwrap(); + assert_eq!(ScalarValue::Binary(None), merged.evaluate().unwrap()); + } } diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index 037ca5b175a..6583f7ba607 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -28,14 +28,14 @@ use crate::aggrs::count_hash::CountHash; use crate::aggrs::vector::VectorFunction as VectorAggrFunction; use crate::function::{Function, FunctionRef}; use crate::function_factory::ScalarFunctionFactory; +#[cfg(feature = "ai_functions")] +use crate::scalars::ai; use crate::scalars::anomaly::AnomalyFunction; use crate::scalars::avg_calc::AvgCalcFunction; use crate::scalars::date::DateFunction; use crate::scalars::expression::ExpressionFunction; use crate::scalars::hll_count::HllCalcFunction; use crate::scalars::ip::IpFunctions; -#[cfg(feature = "ai_functions")] -use crate::scalars::jev::JevFunction; use crate::scalars::json::JsonFunction; use crate::scalars::matches::MatchesFunction; use crate::scalars::matches_term::MatchesTermFunction; @@ -229,7 +229,7 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| MatchesFunction::register(&function_registry); MatchesTermFunction::register(&function_registry); #[cfg(feature = "ai_functions")] - JevFunction::register(&function_registry); + ai::register(&function_registry); // System and administration functions SystemFunction::register(&function_registry); @@ -384,11 +384,17 @@ mod tests { } #[test] - fn test_jev_registration_matches_feature() { - assert_eq!( - FUNCTION_REGISTRY.get_function("jev").is_some(), - cfg!(feature = "ai_functions") - ); + fn test_ai_registration_matches_feature() { + for name in ["ai_match", "ai_choose", "ai_score"] { + assert_eq!( + FUNCTION_REGISTRY.get_function(name).is_some(), + cfg!(feature = "ai_functions"), + "{name}" + ); + } + for name in ["jev", "jev_choice", "jev_score"] { + assert!(FUNCTION_REGISTRY.get_function(name).is_none(), "{name}"); + } } #[test] diff --git a/src/common/function/src/lib.rs b/src/common/function/src/lib.rs index 336f8d489ca..7309dba981d 100644 --- a/src/common/function/src/lib.rs +++ b/src/common/function/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(try_blocks)] - mod admin; mod flush_flow; mod macros; diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 61fb7ec632c..4bc4c17f5bb 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -12,13 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "ai_functions")] +pub(crate) mod ai; pub mod anomaly; pub(crate) mod date; pub mod expression; #[cfg(feature = "geo")] pub mod geo; -#[cfg(feature = "ai_functions")] -pub(crate) mod jev; pub mod json; pub mod matches; pub mod matches_term; diff --git a/src/common/function/src/scalars/ai.rs b/src/common/function/src/scalars/ai.rs new file mode 100644 index 00000000000..f2e17dffd36 --- /dev/null +++ b/src/common/function/src/scalars/ai.rs @@ -0,0 +1,454 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::hash::Hash; +use std::marker::PhantomData; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{Array, StringViewArray, new_null_array}; +use arrow::datatypes::DataType; +use async_trait::async_trait; +use datafusion_common::cast::as_string_view_array; +use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err, not_impl_err}; +use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use futures::{StreamExt, TryStreamExt, stream}; +use reqwest::Client; +use serde::Serialize; +use serde_json::{Map, Value, json}; + +use crate::function_factory::ScalarFunctionFactory; +use crate::function_registry::FunctionRegistry; + +/// Registers the experimental AI matching, classification, and rating functions. +pub(crate) fn register(registry: &FunctionRegistry) { + AiFunction::::register(registry); + AiFunction::::register(registry); + AiFunction::::register(registry); +} + +/// Shared asynchronous execution for AI functions using the Jev backend. +#[derive(PartialEq, Eq, Hash)] +struct AiFunction { + signature: Signature, + question: PhantomData, + enabled: bool, + api_key: Option, + endpoint: String, + model: String, +} + +struct AiRequest<'a, Q: AiQuestion> { + text: &'a str, + prompt: &'a str, + criteria: Arc, +} + +/// Borrows string arguments without expanding constants to the batch size. +enum StringArgument<'a> { + Scalar(Option<&'a str>), + Array(&'a StringViewArray), +} + +impl<'a> StringArgument<'a> { + fn try_new(arg: &'a ColumnarValue) -> Result { + match arg { + ColumnarValue::Scalar(ScalarValue::Utf8View(value)) => { + Ok(Self::Scalar(value.as_deref())) + } + ColumnarValue::Array(array) => Ok(Self::Array(as_string_view_array(array)?)), + _ => exec_err!("AI functions expect Utf8View arguments"), + } + } + + fn value(&self, row: usize) -> Option<&'a str> { + match self { + Self::Scalar(value) => *value, + Self::Array(array) => { + let array = *array; + array.is_valid(row).then(|| array.value(row)) + } + } + } +} + +impl AiFunction { + fn register(registry: &FunctionRegistry) { + registry.register(ScalarFunctionFactory { + name: Q::NAME.to_string(), + factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()), + }); + } + + async fn evaluate(&self, client: &Client, request: AiRequest<'_, Q>) -> Result { + let mut question = json!({ "type": Q::TYPE, "instructions": request.prompt }); + if Q::ARG_COUNT == 3 { + question["criteria"] = json!(request.criteria.as_ref()); + } + let response: Value = client + .post(&self.endpoint) + .json(&json!({ + "model": self.model, + "state": request.text, + "questions": { "matches": question } + })) + .send() + .await + .and_then(|response| response.error_for_status()) + .map_err(|e| exec_datafusion_err!("{} request failed: {e}", Q::NAME))? + .json() + .await + .map_err(|e| exec_datafusion_err!("{} response is not valid JSON: {e}", Q::NAME))?; + + let answer = &response["answers"]["matches"]; + if answer["type"].as_str() != Some(Q::TYPE) { + return exec_err!( + "{} response must contain answers.matches with type {}", + Q::NAME, + Q::TYPE + ); + } + Q::parse_answer(answer, request.criteria.as_ref()) + } + + fn prepare_requests<'a>( + &self, + args: &'a [ColumnarValue], + number_rows: usize, + ) -> Result>>> { + if args.len() != Q::ARG_COUNT { + return exec_err!("{} requires {} arguments", Q::NAME, Q::ARG_COUNT); + } + let args = args + .iter() + .map(StringArgument::try_new) + .collect::>>()?; + let scalar_criteria = args[2..] + .iter() + .all(|arg| matches!(arg, StringArgument::Scalar(_))); + let mut shared_criteria: Option> = None; + + // Validate all non-null rows before making any billable requests. + (0..number_rows) + .map(|row| { + let values: Option> = args.iter().map(|arg| arg.value(row)).collect(); + values + .map(|values| { + // Initialize lazily so NULL rows never validate otherwise invalid criteria. + let criteria = match &shared_criteria { + Some(criteria) => Arc::clone(criteria), + None => { + let criteria = Arc::new(Q::parse_criteria(&values[2..])?); + if scalar_criteria { + shared_criteria = Some(Arc::clone(&criteria)); + } + criteria + } + }; + Ok(AiRequest { + text: values[0], + prompt: values[1], + criteria, + }) + }) + .transpose() + }) + .collect() + } + + fn client(&self) -> Result { + if !self.enabled { + return exec_err!( + "{} is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it", + Q::NAME + ); + } + let key = self + .api_key + .as_deref() + .filter(|key| !key.trim().is_empty()) + .ok_or_else(|| { + exec_datafusion_err!("{} requires the JEV_API_KEY environment variable", Q::NAME) + })?; + let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}")) + .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?; + authorization.set_sensitive(true); + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(reqwest::header::AUTHORIZATION, authorization); + Client::builder() + .default_headers(headers) + .timeout(Duration::from_secs(30)) + .build() + .map_err(|e| exec_datafusion_err!("failed to create {} HTTP client: {e}", Q::NAME)) + } +} + +impl Default for AiFunction { + fn default() -> Self { + Self { + // External model evaluations must not be constant-folded during planning. + signature: Signature::exact( + vec![DataType::Utf8View; Q::ARG_COUNT], + Volatility::Volatile, + ), + question: PhantomData, + enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"), + api_key: std::env::var("JEV_API_KEY").ok(), + endpoint: std::env::var("JEV_ENDPOINT") + .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()), + model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()), + } + } +} + +impl fmt::Debug for AiFunction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Credentials must never appear in plans or diagnostic output. + f.debug_struct("AiFunction") + .field("name", &Q::NAME) + .field("signature", &self.signature) + .field("enabled", &self.enabled) + .field("model", &self.model) + .finish_non_exhaustive() + } +} + +impl ScalarUDFImpl for AiFunction { + fn name(&self) -> &str { + Q::NAME + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(Q::return_type()) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + not_impl_err!("{} can only be called from async contexts", Q::NAME) + } +} + +#[async_trait] +impl AsyncScalarUDFImpl for AiFunction { + async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { + let rows = self.prepare_requests(&args.args, args.number_rows)?; + + if rows.iter().all(Option::is_none) { + return Ok(ColumnarValue::Array(new_null_array( + &Q::return_type(), + args.number_rows, + ))); + } + + // Reuse connections within the batch; bound in-flight requests and preserve row order. + let client = self.client()?; + // Collect futures to avoid async_trait's higher-ranked lifetime inference issue. + let requests: Vec<_> = rows + .into_iter() + .map(|row| { + let client = &client; + async move { + match row { + Some(request) => self.evaluate(client, request).await, + None => ScalarValue::try_from(&Q::return_type()), + } + } + }) + .collect(); + // This limit is per expression/batch invocation, not per query or process. + // Concurrent partitions and queries can each have their own in-flight requests. + let answers: Vec = stream::iter(requests).buffered(8).try_collect().await?; + Ok(ColumnarValue::Array(ScalarValue::iter_to_array(answers)?)) + } +} + +/// The request criteria and scalar answer contract for an AI question type. +trait AiQuestion: fmt::Debug + Eq + Hash + Send + Sync + 'static { + type Criteria: Serialize + Send + Sync; + + const NAME: &'static str; + const TYPE: &'static str; + const ARG_COUNT: usize; + + fn return_type() -> DataType; + fn parse_criteria(args: &[&str]) -> Result; + fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result; +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct Noul; + +impl AiQuestion for Noul { + type Criteria = (); + + const NAME: &'static str = "ai_match"; + const TYPE: &'static str = "noul"; + const ARG_COUNT: usize = 2; + + fn return_type() -> DataType { + DataType::Float64 + } + + fn parse_criteria(_args: &[&str]) -> Result<()> { + Ok(()) + } + + fn parse_answer(answer: &Value, _criteria: &()) -> Result { + numeric_answer(Self::NAME, &answer["noul"], "noul", 1.0) + .map(|probability| ScalarValue::Float64(Some(probability))) + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct Choice; + +impl AiQuestion for Choice { + type Criteria = Map; + + const NAME: &'static str = "ai_choose"; + const TYPE: &'static str = "choice"; + const ARG_COUNT: usize = 3; + + fn return_type() -> DataType { + DataType::Utf8 + } + + fn parse_criteria(args: &[&str]) -> Result { + let criteria: Value = serde_json::from_str(args[0]) + .map_err(|e| exec_datafusion_err!("ai_choose criteria is not valid JSON: {e}"))?; + match criteria { + Value::Object(options) + if (1..=255).contains(&options.len()) + && options.values().all(|v| v.is_null() || is_description(v)) => + { + Ok(options) + } + _ => exec_err!( + "ai_choose criteria must be a JSON object with 1 to 255 options; descriptions must be strings, objects, arrays, or null" + ), + } + } + + fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result { + let choice = answer["choice"] + .as_str() + .filter(|choice| criteria.contains_key(*choice)) + .ok_or_else(|| { + exec_datafusion_err!("ai_choose response must contain a choice from the criteria") + })?; + Ok(ScalarValue::Utf8(Some(choice.to_string()))) + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct Score; + +impl AiQuestion for Score { + type Criteria = Vec; + + const NAME: &'static str = "ai_score"; + const TYPE: &'static str = "score"; + const ARG_COUNT: usize = 3; + + fn return_type() -> DataType { + DataType::BinaryView + } + + fn parse_criteria(args: &[&str]) -> Result { + let criteria: Value = serde_json::from_str(args[0]) + .map_err(|e| exec_datafusion_err!("ai_score criteria is not valid JSON: {e}"))?; + match criteria { + Value::Array(levels) + if (2..=10).contains(&levels.len()) && levels.iter().all(is_description) => + { + Ok(levels) + } + _ => exec_err!( + "ai_score criteria must be a JSON array with 2 to 10 levels; descriptions must be strings, objects, or arrays" + ), + } + } + + fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result { + let score = numeric_answer( + Self::NAME, + &answer["score"], + "score", + (criteria.len() - 1) as f64, + )?; + let confidence = numeric_answer(Self::NAME, &answer["confidence"], "confidence", 1.0)?; + let probabilities = score_probabilities(&answer["probabilities"], criteria.len())?; + let object = jsonb::Object::from([ + ("score".to_string(), jsonb::Value::from(score)), + ("confidence".to_string(), jsonb::Value::from(confidence)), + ( + "probabilities".to_string(), + jsonb::Value::Array(probabilities.into_iter().map(jsonb::Value::from).collect()), + ), + ]); + Ok(ScalarValue::BinaryView(Some( + jsonb::Value::Object(object).to_vec(), + ))) + } +} + +fn score_probabilities(distribution: &Value, level_count: usize) -> Result> { + if distribution + .as_object() + .is_none_or(|probabilities| probabilities.len() != level_count) + { + return exec_err!( + "ai_score response must contain probabilities for all {level_count} levels" + ); + } + let probabilities = (0..level_count) + .map(|level| { + numeric_answer( + Score::NAME, + &distribution[level.to_string()], + "probability", + 1.0, + ) + }) + .collect::>>()?; + // Allow small rounding differences without renormalizing the provider's distribution. + if (probabilities.iter().sum::() - 1.0).abs() > 1e-6 { + return exec_err!("ai_score response probabilities must sum to 1 within 1e-6"); + } + Ok(probabilities) +} + +fn is_description(description: &Value) -> bool { + matches!( + description, + Value::String(_) | Value::Object(_) | Value::Array(_) + ) +} + +fn numeric_answer(name: &str, answer: &Value, field: &str, max: f64) -> Result { + answer + .as_f64() + .filter(|score| (0.0..=max).contains(score)) + .ok_or_else(|| { + exec_datafusion_err!("{name} response must contain a finite {field} in [0, {max}]") + }) +} diff --git a/src/common/function/src/scalars/ai/tests.rs b/src/common/function/src/scalars/ai/tests.rs new file mode 100644 index 00000000000..0b001e22e7a --- /dev/null +++ b/src/common/function/src/scalars/ai/tests.rs @@ -0,0 +1,766 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +use axum::http::{HeaderMap, StatusCode}; +use axum::routing::post; +use axum::{Json, Router}; +use datafusion::assert_batches_eq; +use datafusion::prelude::{SessionConfig, SessionContext}; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; + +use super::*; +use crate::function::FunctionContext; +use crate::function_registry::FUNCTION_REGISTRY; + +struct MockServer { + endpoint: String, + task: JoinHandle<()>, +} + +impl MockServer { + async fn start(router: Router) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + Self { endpoint, task } + } + + fn function(&self) -> AiFunction { + AiFunction { + enabled: true, + api_key: Some("test-key".to_string()), + endpoint: self.endpoint.clone(), + model: "test-model".to_string(), + ..Default::default() + } + } +} + +impl Drop for MockServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +fn context(function: AiFunction) -> SessionContext { + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf()); + let json_get_float = FUNCTION_REGISTRY.get_function("json_get_float").unwrap(); + ctx.register_udf(json_get_float.provide(FunctionContext::default())); + ctx +} + +#[test] +fn test_ai_scalar_criteria_share_allocation_and_skip_null_rows() { + fn assert_shared_criteria(criteria: &str) { + let function = AiFunction::::default(); + for (texts, criteria) in [ + (vec![None, Some("first"), Some("second")], criteria), + (vec![None, None], "not json"), + (vec![], "not json"), + ] { + let number_rows = texts.len(); + let expected_requests = texts.iter().flatten().count(); + let args = vec![ + ColumnarValue::Array(Arc::new(StringViewArray::from(texts))), + ColumnarValue::Scalar(ScalarValue::Utf8View(Some("prompt".to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8View(Some(criteria.to_string()))), + ]; + let requests = function.prepare_requests(&args, number_rows).unwrap(); + assert_eq!(requests.len(), number_rows); + assert_eq!(requests.iter().flatten().count(), expected_requests); + + // Regression: parsed constant criteria must not be duplicated for each row. + let mut criteria = requests.iter().flatten().map(|request| &request.criteria); + if let Some(first) = criteria.next() { + assert!(criteria.all(|other| Arc::ptr_eq(first, other))); + } + } + } + + assert_shared_criteria::(r#"{"billing":null,"technical":"errors"}"#); + assert_shared_criteria::(r#"["low","high"]"#); +} + +#[tokio::test] +async fn test_ai_match_sql_returns_and_filters_noul_probability() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post( + |headers: HeaderMap, Json(request): Json| async move { + assert_eq!(headers["authorization"], "Bearer test-key"); + assert_eq!(request["model"], "test-model"); + assert_eq!(request["questions"]["matches"]["type"], "noul"); + // The prompt is passed through verbatim. + assert_eq!( + request["questions"]["matches"]["instructions"], + "payment failed" + ); + let probability = match request["state"].as_str().unwrap() { + "failed" => { + // Complete this row after later rows to exercise result alignment. + tokio::time::sleep(Duration::from_millis(20)).await; + 0.95 + } + "boundary" => 0.8, + "recovered" => 0.1, + "impossible" => 0.0, + "certain" => 1.0, + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}})) + }, + ), + )) + .await; + let ctx = context(server.function::()); + let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL), (5, 'impossible'), (6, 'certain')) AS events(id, message)"; + let batches = ctx + .sql(&format!( + "SELECT id FROM {events} WHERE ai_match((message), 'payment failed') >= 0.8 ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 6 |", "+----+" + ], + &batches + ); + + let batches = ctx + .sql(&format!( + "SELECT id, ai_match(message, 'payment failed') AS score FROM {events} ORDER BY score DESC NULLS LAST" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Float64); + assert_batches_eq!( + [ + "+----+-------+", + "| id | score |", + "+----+-------+", + "| 6 | 1.0 |", + "| 1 | 0.95 |", + "| 2 | 0.8 |", + "| 3 | 0.1 |", + "| 5 | 0.0 |", + "| 4 | |", + "+----+-------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_ai_match_nulls_need_no_api_key() { + let ctx = context(AiFunction:: { + enabled: true, + api_key: None, + ..Default::default() + }); + let batches = ctx + .sql("SELECT ai_match(NULL, 'condition') AS a, ai_match('text', NULL) AS b") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(0).data_type(), &DataType::Float64); + assert_batches_eq!( + [ + "+---+---+", + "| a | b |", + "+---+---+", + "| | |", + "+---+---+" + ], + &batches + ); + + let error = ctx + .sql("SELECT ai_match('text', 'condition')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains("JEV_API_KEY"), "{error}"); + + let ctx = context(AiFunction:: { + enabled: false, + api_key: None, + ..Default::default() + }); + let error = ctx + .sql("SELECT ai_match('text', 'condition')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"), + "{error}" + ); +} + +#[tokio::test] +async fn test_ai_match_bad_api_responses_fail_the_query() { + for (status, body, expected) in [ + ( + StatusCode::TOO_MANY_REQUESTS, + "{}", + "ai_match request failed", + ), + ( + StatusCode::OK, + "not json", + "ai_match response is not valid JSON", + ), + ( + StatusCode::OK, + r#"{"answers":{}}"#, + "ai_match response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#, + "ai_match response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"noul","noul":-0.1}}}"#, + "ai_match response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#, + "ai_match response must contain", + ), + ] { + let server = MockServer::start( + Router::new().route("/v1/systemone", post(move || async move { (status, body) })), + ) + .await; + let ctx = context(server.function::()); + let error = ctx + .sql("SELECT ai_match('text', 'condition')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + } +} + +#[tokio::test] +async fn test_ai_choose_returns_and_filters_supplied_labels() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(|Json(request): Json| async move { + let question = &request["questions"]["matches"]; + assert_eq!(question["type"], "choice"); + assert_eq!(question["instructions"], "Route the ticket"); + let choice = match request["state"].as_str().unwrap() { + "refund" => { + assert_eq!(question["criteria"], json!({"billing": null, "technical": {"examples": ["crash"]}})); + tokio::time::sleep(Duration::from_millis(20)).await; + "billing" + } + "crash" => { + assert_eq!(question["criteria"], json!({"billing": ["refunds"], "technical": "errors"})); + "technical" + } + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "choice", "choice": choice, "confidence": 0.7}}})) + }), + )) + .await; + let ctx = context(server.function::()); + let events = r#"(VALUES + (1, 'refund', 'Route the ticket', '{"billing":null,"technical":{"examples":["crash"]}}'), + (2, 'crash', 'Route the ticket', '{"billing":["refunds"],"technical":"errors"}'), + (3, NULL, 'Route the ticket', 'invalid JSON'), + (4, 'refund', NULL, 'invalid JSON'), + (5, 'refund', 'Route the ticket', NULL) + ) AS events(id, message, prompt, criteria)"#; + let batches = ctx + .sql(&format!( + "SELECT id, ai_choose(message, prompt, criteria) AS team FROM {events} ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Utf8); + assert_batches_eq!( + [ + "+----+-----------+", + "| id | team |", + "+----+-----------+", + "| 1 | billing |", + "| 2 | technical |", + "| 3 | |", + "| 4 | |", + "| 5 | |", + "+----+-----------+" + ], + &batches + ); + let batches = ctx + .sql(&format!( + "SELECT id FROM {events} WHERE ai_choose(message, prompt, criteria) = 'billing'" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); + + let batches = ctx + .sql( + r#"SELECT id, ai_choose(message, 'Route the ticket', + '{"billing":null,"technical":{"examples":["crash"]}}') AS team + FROM (VALUES (1, 'refund'), (2, NULL), (3, 'refund')) AS events(id, message) + ORDER BY id"#, + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+---------+", + "| id | team |", + "+----+---------+", + "| 1 | billing |", + "| 2 | |", + "| 3 | billing |", + "+----+---------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_ai_score_returns_fractional_scores_on_each_rows_scale() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(|Json(request): Json| async move { + let question = &request["questions"]["matches"]; + assert_eq!(question["type"], "score"); + assert_eq!(question["instructions"], "Rate severity"); + let (score, probabilities) = match request["state"].as_str().unwrap() { + "low" => { + assert_eq!(question["criteria"], json!(["low", {"description": "medium"}, ["high"]])); + tokio::time::sleep(Duration::from_millis(20)).await; + (0.0, json!({"0":1.0,"1":0.0,"2":0.0})) + } + "fractional" => (1.25, json!({"2":0.25,"0":0.0,"1":0.75})), + "high" => (2.0, json!({"0":0.0,"1":0.0,"2":1.0})), + "binary" => { + assert_eq!(question["criteria"], json!(["low", "high"])); + (1.0, json!({"0":0.0,"1":1.0})) + } + "ten levels" => { + assert_eq!(question["criteria"].as_array().unwrap().len(), 10); + let probabilities: Map = (0..10) + .map(|level| (level.to_string(), json!(if level == 9 { 1.0 } else { 0.0 }))) + .collect(); + (9.0, Value::Object(probabilities)) + } + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "score", "score": score, "confidence": 0.4, "probabilities": probabilities}}})) + }), + )) + .await; + let ctx = context(server.function::()); + let batches = ctx + .sql( + r#" + SELECT id, json_get_float(ai_score(message, 'Rate severity', criteria), 'score') AS severity + FROM (VALUES + (1, 'low', '["low", {"description":"medium"}, ["high"]]'), + (2, 'fractional', '["low", "medium", "high"]'), + (3, 'high', '["low", "medium", "high"]'), + (4, 'binary', '["low", "high"]'), + (5, 'ten levels', '["a","b","c","d","e","f","g","h","i","j"]'), + (6, NULL, 'invalid JSON'), + (7, 'high', NULL) + ) AS events(id, message, criteria) + ORDER BY severity DESC NULLS LAST, id + "#, + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Float64); + assert_batches_eq!( + [ + "+----+----------+", + "| id | severity |", + "+----+----------+", + "| 5 | 9.0 |", + "| 3 | 2.0 |", + "| 2 | 1.25 |", + "| 4 | 1.0 |", + "| 1 | 0.0 |", + "| 6 | |", + "| 7 | |", + "+----+----------+" + ], + &batches + ); + + let batches = ctx + .sql( + r#"SELECT id, json_get_float(ai_score(message, 'Rate severity', + '["low",{"description":"medium"},["high"]]'), 'score') AS severity + FROM (VALUES (1, 'low'), (2, 'fractional'), (3, 'high'), (4, NULL)) AS events(id, message) + ORDER BY id"#, + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+----------+", + "| id | severity |", + "+----+----------+", + "| 1 | 0.0 |", + "| 2 | 1.25 |", + "| 3 | 2.0 |", + "| 4 | |", + "+----+----------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_ai_score_preserves_uncertainty_in_one_evaluation() { + let calls = Arc::new(AtomicUsize::new(0)); + let request_calls = calls.clone(); + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(move |Json(request): Json| { + request_calls.fetch_add(1, Ordering::Relaxed); + async move { + let (confidence, probabilities) = match request["state"].as_str().unwrap() { + "medium" => (1.0, json!({"2":0.0,"0":0.0,"1":1.0})), + "split" => (0.25, json!({"2":0.5,"0":0.5,"1":0.0})), + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers":{"matches":{ + "type":"score", "score":1.0, "confidence":confidence, + "probabilities":probabilities + }}})) + } + }), + )) + .await; + let ctx = context(server.function::()); + let rated = r#"(SELECT id, ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating + FROM (VALUES (1, 'medium'), (2, 'split'), (3, NULL)) AS events(id, message)) AS rated"#; + let batches = ctx + .sql(&format!( + "SELECT id, json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[0]') AS p0, + json_get_float(rating, 'probabilities[1]') AS p1, + json_get_float(rating, 'probabilities[2]') AS p2 + FROM {rated} ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+-------+------------+-----+-----+-----+", + "| id | score | confidence | p0 | p1 | p2 |", + "+----+-------+------------+-----+-----+-----+", + "| 1 | 1.0 | 1.0 | 0.0 | 1.0 | 0.0 |", + "| 2 | 1.0 | 0.25 | 0.5 | 0.0 | 0.5 |", + "| 3 | | | | | |", + "+----+-------+------------+-----+-----+-----+" + ], + &batches + ); + assert_eq!(calls.load(Ordering::Relaxed), 2); + + let batches = ctx + .sql(&format!( + "SELECT id FROM {rated} WHERE json_get_float(rating, 'confidence') >= 0.8 + ORDER BY json_get_float(rating, 'score') DESC" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); + assert_eq!(calls.load(Ordering::Relaxed), 4); +} + +#[test] +fn test_ai_score_preserves_provider_precision() { + let answer = json!({ + "score":1.6, "confidence":0.92, + "probabilities":{"2":0.6999995,"0":0.1,"1":0.2} + }); + let criteria = vec![json!("low"), json!("medium"), json!("high")]; + let ScalarValue::BinaryView(Some(bytes)) = Score::parse_answer(&answer, &criteria).unwrap() + else { + panic!("expected a JSONB object"); + }; + let object: Value = + serde_json::from_str(&jsonb::from_slice(&bytes).unwrap().to_string()).unwrap(); + assert_eq!(object.as_object().unwrap().len(), 3); + assert_eq!(object["score"].as_f64(), Some(1.6)); + assert_eq!(object["confidence"].as_f64(), Some(0.92)); + let probabilities: Vec = object["probabilities"] + .as_array() + .unwrap() + .iter() + .map(|probability| probability.as_f64().unwrap()) + .collect(); + assert_eq!(probabilities, vec![0.1, 0.2, 0.6999995]); +} + +#[tokio::test] +async fn test_ai_score_rejects_invalid_uncertainty() { + for (field, value) in [ + ("confidence", Value::Null), + ("confidence", json!(-0.1)), + ("confidence", json!(1.1)), + ("confidence", json!("0.9")), + ("probabilities", Value::Null), + ("probabilities", json!([0.0, 1.0, 0.0])), + ("probabilities", json!({"0":1.0})), + ("probabilities", json!({"0":0.5,"2":0.5,"3":0.0})), + ("probabilities", json!({"0":0.0,"1":1.0,"2":0.0,"3":0.0})), + ("probabilities", json!({"0":-0.1,"1":0.6,"2":0.5})), + ("probabilities", json!({"0":0.0,"1":1.1,"2":0.0})), + ("probabilities", json!({"0":0.0,"1":0.95,"2":"0.05"})), + ("probabilities", json!({"0":0.0,"1":0.0,"2":0.0})), + ("probabilities", json!({"0":0.5,"1":0.5,"2":0.5})), + ] { + let mut answer = json!({ + "type":"score", "score":1.0, "confidence":0.9, + "probabilities":{"0":0.0,"1":1.0,"2":0.0} + }); + answer[field] = value; + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(move || { + let answer = answer.clone(); + async move { Json(json!({"answers":{"matches":answer}})) } + }), + )) + .await; + let ctx = context(server.function::()); + let error = ctx + .sql(r#"SELECT ai_score('text', 'prompt', '["low","medium","high"]')"#) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains("ai_score response"), "{error}"); + } +} + +#[tokio::test] +async fn test_ai_invalid_criteria_fail_before_requests() { + let choice_ctx = context(AiFunction:: { + enabled: true, + api_key: None, + ..Default::default() + }); + let score_ctx = context(AiFunction:: { + enabled: true, + api_key: None, + ..Default::default() + }); + let too_many_choices: Map = (0..256) + .map(|index| (index.to_string(), Value::Null)) + .collect(); + for (ctx, name, valid, invalid) in [ + ( + &choice_ctx, + "ai_choose", + r#"{"billing":null}"#, + vec![ + "not json".to_string(), + "null".to_string(), + "{}".to_string(), + "[]".to_string(), + r#"{"billing":1}"#.to_string(), + r#"{"billing":true}"#.to_string(), + Value::Object(too_many_choices).to_string(), + ], + ), + ( + &score_ctx, + "ai_score", + r#"["low","high"]"#, + vec![ + "not json".to_string(), + "null".to_string(), + "{}".to_string(), + "[]".to_string(), + r#"["low"]"#.to_string(), + r#"["low",null]"#.to_string(), + r#"["low",1]"#.to_string(), + r#"["low",true]"#.to_string(), + json!(vec!["level"; 11]).to_string(), + ], + ), + ] { + let error = ctx + .sql(&format!( + "SELECT {name}(message, 'prompt', 'not json') FROM (VALUES (NULL), ('text')) AS input(message)" + )) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error.to_string().contains(&format!("{name} criteria")), + "{error}" + ); + for criteria in invalid { + // A later invalid row must fail before even creating a client for the first row. + let error = ctx.sql(&format!( + "SELECT {name}('text', 'prompt', criteria) FROM (VALUES ('{valid}'), ('{criteria}')) AS input(criteria)" + )).await.unwrap().collect().await.unwrap_err(); + assert!( + error.to_string().contains(&format!("{name} criteria")), + "{error}" + ); + } + } +} + +#[tokio::test] +async fn test_ai_choose_and_score_reject_invalid_answers() { + for (name, criteria, answer) in [ + ( + "ai_choose", + r#"{"billing":null}"#, + json!({"type":"score","choice":"billing"}), + ), + ( + "ai_choose", + r#"{"billing":null}"#, + json!({"type":"choice","choice":"unknown"}), + ), + ( + "ai_choose", + r#"{"billing":null}"#, + json!({"type":"choice","choice":1}), + ), + ("ai_choose", r#"{"billing":null}"#, json!({"type":"choice"})), + ( + "ai_score", + r#"["low","medium","high"]"#, + json!({"type":"noul","score":0.5}), + ), + ( + "ai_score", + r#"["low","medium","high"]"#, + json!({"type":"score","score":2.1}), + ), + ( + "ai_score", + r#"["low","high"]"#, + json!({"type":"score","score":-0.1}), + ), + ( + "ai_score", + r#"["low","high"]"#, + json!({"type":"score","score":"NaN"}), + ), + ("ai_score", r#"["low","high"]"#, json!({"type":"score"})), + ] { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(move || { + let answer = answer.clone(); + async move { Json(json!({"answers":{"matches":answer}})) } + }), + )) + .await; + let ctx = context(server.function::()); + ctx.register_udf( + AsyncScalarUDF::new(Arc::new(server.function::())).into_scalar_udf(), + ); + let error = ctx + .sql(&format!("SELECT {name}('text', 'prompt', '{criteria}')")) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains(&format!("{name} response must contain")), + "{error}" + ); + } +} + +#[tokio::test] +#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"] +async fn test_ai_match_live() { + let ctx = SessionContext::new(); + let function = FUNCTION_REGISTRY.get_function("ai_match").unwrap(); + ctx.register_udf(function.provide(FunctionContext::default())); + let batches = ctx.sql( + "SELECT id FROM (VALUES + (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'), + (2, 'Payment succeeded on the second retry; the payment is complete.'), + (3, 'User logged in successfully.') + ) AS events(id, message) + WHERE ai_match(message, 'The event reports that a payment still failed after retries.') >= 0.8 + ORDER BY id", + ).await.unwrap().collect().await.unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); +} diff --git a/src/common/function/src/scalars/jev.rs b/src/common/function/src/scalars/jev.rs deleted file mode 100644 index 1c6475e6fac..00000000000 --- a/src/common/function/src/scalars/jev.rs +++ /dev/null @@ -1,209 +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. - -#[cfg(test)] -mod tests; - -use std::fmt; -use std::sync::Arc; -use std::time::Duration; - -use arrow::array::BooleanArray; -use arrow::datatypes::DataType; -use async_trait::async_trait; -use datafusion_common::cast::{as_float64_array, as_string_view_array}; -use datafusion_common::utils::take_function_args; -use datafusion_common::{Result, exec_datafusion_err, exec_err, not_impl_err}; -use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; -use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; -use futures::{StreamExt, TryStreamExt, stream}; -use reqwest::Client; -use serde_json::{Value, json}; - -use crate::function_factory::ScalarFunctionFactory; -use crate::function_registry::FunctionRegistry; - -/// Experimental natural-language predicate: `jev(text, statement, threshold)`. -/// Each non-null row is a Noul question; its probability is compared with `>=`. -#[derive(PartialEq, Eq, Hash)] -pub(crate) struct JevFunction { - signature: Signature, - enabled: bool, - api_key: Option, - endpoint: String, - model: String, -} - -impl JevFunction { - pub(crate) fn register(registry: &FunctionRegistry) { - registry.register(ScalarFunctionFactory { - name: "jev".to_string(), - factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()), - }); - } - - async fn evaluate(&self, client: &Client, text: &str, statement: &str) -> Result { - let response: Value = client - .post(&self.endpoint) - .json(&json!({ - "model": self.model, - "state": text, - "questions": { - "matches": { "type": "noul", "instructions": statement } - } - })) - .send() - .await - .and_then(|response| response.error_for_status()) - .map_err(|e| exec_datafusion_err!("jev request failed: {e}"))? - .json() - .await - .map_err(|e| exec_datafusion_err!("jev response is not valid JSON: {e}"))?; - - let answer = &response["answers"]["matches"]; - let probability = answer["noul"].as_f64().filter(|p| (0.0..=1.0).contains(p)); - match (answer["type"].as_str(), probability) { - (Some("noul"), Some(probability)) => Ok(probability), - _ => exec_err!( - "jev response must contain answers.matches with type noul and a probability in [0, 1]" - ), - } - } - - fn client(&self) -> Result { - if !self.enabled { - return exec_err!( - "jev is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it" - ); - } - let key = self - .api_key - .as_deref() - .filter(|key| !key.trim().is_empty()) - .ok_or_else(|| { - exec_datafusion_err!("jev requires the JEV_API_KEY environment variable") - })?; - let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}")) - .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?; - authorization.set_sensitive(true); - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert(reqwest::header::AUTHORIZATION, authorization); - Client::builder() - .default_headers(headers) - .timeout(Duration::from_secs(30)) - .build() - .map_err(|e| exec_datafusion_err!("failed to create jev HTTP client: {e}")) - } -} - -impl Default for JevFunction { - fn default() -> Self { - Self { - // External model evaluations must not be constant-folded during planning. - signature: Signature::exact( - vec![DataType::Utf8View, DataType::Utf8View, DataType::Float64], - Volatility::Volatile, - ), - enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"), - api_key: std::env::var("JEV_API_KEY").ok(), - endpoint: std::env::var("JEV_ENDPOINT") - .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()), - model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()), - } - } -} - -impl fmt::Debug for JevFunction { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Credentials must never appear in plans or diagnostic output. - f.debug_struct("JevFunction") - .field("signature", &self.signature) - .field("enabled", &self.enabled) - .field("model", &self.model) - .finish_non_exhaustive() - } -} - -impl ScalarUDFImpl for JevFunction { - fn name(&self) -> &str { - "jev" - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Boolean) - } - - fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { - not_impl_err!("jev can only be called from async contexts") - } -} - -#[async_trait] -impl AsyncScalarUDFImpl for JevFunction { - async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { - let arrays = args - .args - .into_iter() - .map(|arg| arg.into_array(args.number_rows)) - .collect::>>()?; - let [texts, statements, thresholds] = take_function_args(self.name(), arrays)?; - let texts = as_string_view_array(&texts)?; - let statements = as_string_view_array(&statements)?; - let thresholds = as_float64_array(&thresholds)?; - let rows: Vec<_> = texts - .iter() - .zip(statements.iter()) - .zip(thresholds.iter()) - .map(|((text, statement), threshold)| Some((text?, statement?, threshold?))) - .collect(); - - // Validate the whole batch before making any billable requests. - for (_, _, threshold) in rows.iter().flatten() { - if !(0.0..=1.0).contains(threshold) { - return exec_err!("jev threshold must be finite and in [0, 1]"); - } - } - if rows.iter().all(Option::is_none) { - return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( - args.number_rows, - )))); - } - - // Reuse connections within the batch; bound in-flight requests and preserve row order. - let client = self.client()?; - let requests: Vec<_> = rows - .into_iter() - .map(|row| { - let client = &client; - async move { - match row { - Some((text, statement, threshold)) => self - .evaluate(client, text, statement) - .await - .map(|p| Some(p >= threshold)), - None => Ok(None), - } - } - }) - .collect(); - // This limit is per expression/batch invocation, not per query or process. - // Concurrent partitions and queries can each have their own in-flight requests. - let matches: Vec> = stream::iter(requests).buffered(8).try_collect().await?; - Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(matches)))) - } -} diff --git a/src/common/function/src/scalars/jev/tests.rs b/src/common/function/src/scalars/jev/tests.rs deleted file mode 100644 index 82d7ab03013..00000000000 --- a/src/common/function/src/scalars/jev/tests.rs +++ /dev/null @@ -1,250 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use axum::http::{HeaderMap, StatusCode}; -use axum::routing::post; -use axum::{Json, Router}; -use datafusion::assert_batches_eq; -use datafusion::prelude::{SessionConfig, SessionContext}; -use tokio::net::TcpListener; -use tokio::task::JoinHandle; - -use super::*; -use crate::function::FunctionContext; -use crate::function_registry::FUNCTION_REGISTRY; - -struct MockServer { - endpoint: String, - task: JoinHandle<()>, -} - -impl MockServer { - async fn start(router: Router) -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); - Self { endpoint, task } - } - - fn function(&self) -> JevFunction { - JevFunction { - enabled: true, - api_key: Some("test-key".to_string()), - endpoint: self.endpoint.clone(), - model: "test-model".to_string(), - ..Default::default() - } - } -} - -impl Drop for MockServer { - fn drop(&mut self) { - self.task.abort(); - } -} - -fn context(function: JevFunction) -> SessionContext { - let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); - ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf()); - ctx -} - -#[tokio::test] -async fn test_jev_sql_filters_using_noul_probability() { - let server = MockServer::start(Router::new().route( - "/v1/systemone", - post( - |headers: HeaderMap, Json(request): Json| async move { - assert_eq!(headers["authorization"], "Bearer test-key"); - assert_eq!(request["model"], "test-model"); - assert_eq!(request["questions"]["matches"]["type"], "noul"); - // The SQL statement is passed through verbatim, not rewritten as a prompt. - assert_eq!( - request["questions"]["matches"]["instructions"], - "payment failed" - ); - let probability = match request["state"].as_str().unwrap() { - "failed" => { - // Complete this row after later rows to exercise result alignment. - tokio::time::sleep(Duration::from_millis(20)).await; - 0.95 - } - "boundary" => 0.8, - "recovered" => 0.1, - unexpected => panic!("unexpected state: {unexpected}"), - }; - Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}})) - }, - ), - )) - .await; - let ctx = context(server.function()); - let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL)) AS events(id, message)"; - let batches = ctx - .sql(&format!( - "SELECT id FROM {events} WHERE jev((message), 'payment failed', 0.8) ORDER BY id" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_batches_eq!( - ["+----+", "| id |", "+----+", "| 1 |", "| 2 |", "+----+"], - &batches - ); - - let batches = ctx - .sql(&format!( - "SELECT id, jev(message, 'payment failed', 0.8) AS matched FROM {events} ORDER BY id" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_batches_eq!( - [ - "+----+---------+", - "| id | matched |", - "+----+---------+", - "| 1 | true |", - "| 2 | true |", - "| 3 | false |", - "| 4 | |", - "+----+---------+" - ], - &batches - ); -} - -#[tokio::test] -async fn test_jev_nulls_and_invalid_thresholds_need_no_api_key() { - let ctx = context(JevFunction { - enabled: true, - api_key: None, - ..Default::default() - }); - let batches = ctx - .sql("SELECT jev(NULL, 'condition', 0.8) AS a, jev('text', NULL, 0.8) AS b, jev('text', 'condition', NULL) AS c") - .await.unwrap().collect().await.unwrap(); - assert_batches_eq!( - [ - "+---+---+---+", - "| a | b | c |", - "+---+---+---+", - "| | | |", - "+---+---+---+" - ], - &batches - ); - - for threshold in ["-0.1", "1.1", "'NaN'::DOUBLE", "'Infinity'::DOUBLE"] { - let error = ctx - .sql(&format!("SELECT jev('text', 'condition', {threshold})")) - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!( - error - .to_string() - .contains("jev threshold must be finite and in [0, 1]"), - "{error}" - ); - } - let error = ctx - .sql("SELECT jev('text', 'condition', 0.8)") - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!(error.to_string().contains("JEV_API_KEY"), "{error}"); - - let ctx = context(JevFunction { - enabled: false, - api_key: None, - ..Default::default() - }); - let error = ctx - .sql("SELECT jev('text', 'condition', 0.8)") - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!( - error - .to_string() - .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"), - "{error}" - ); -} - -#[tokio::test] -async fn test_jev_bad_api_responses_fail_the_query() { - for (status, body, expected) in [ - (StatusCode::TOO_MANY_REQUESTS, "{}", "jev request failed"), - (StatusCode::OK, "not json", "jev response is not valid JSON"), - ( - StatusCode::OK, - r#"{"answers":{}}"#, - "jev response must contain", - ), - ( - StatusCode::OK, - r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#, - "jev response must contain", - ), - ( - StatusCode::OK, - r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#, - "jev response must contain", - ), - ] { - let server = MockServer::start( - Router::new().route("/v1/systemone", post(move || async move { (status, body) })), - ) - .await; - let ctx = context(server.function()); - let error = ctx - .sql("SELECT jev('text', 'condition', 0.8)") - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!(error.to_string().contains(expected), "{error}"); - } -} - -#[tokio::test] -#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"] -async fn test_jev_live() { - let ctx = SessionContext::new(); - let function = FUNCTION_REGISTRY.get_function("jev").unwrap(); - ctx.register_udf(function.provide(FunctionContext::default())); - let batches = ctx.sql( - "SELECT id FROM (VALUES - (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'), - (2, 'Payment succeeded on the second retry; the payment is complete.'), - (3, 'User logged in successfully.') - ) AS events(id, message) - WHERE jev(message, 'The event reports that a payment still failed after retries.', 0.8) - ORDER BY id", - ).await.unwrap().collect().await.unwrap(); - assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); -} diff --git a/src/common/function/src/scalars/matches.rs b/src/common/function/src/scalars/matches.rs index 14889957ee9..109ad1daeb2 100644 --- a/src/common/function/src/scalars/matches.rs +++ b/src/common/function/src/scalars/matches.rs @@ -1119,11 +1119,12 @@ mod test { ]; for (query, expected) in cases { - let result: Result<()> = try { + let result: Result<()> = (|| { let parser = ParserContext { stack: vec![] }; let ast = parser.parse_pattern(query)?; let _ast = ast.transform_ast()?; - }; + Ok(()) + })(); assert!(result.is_err(), "{query}"); let actual_error = result.unwrap_err().to_string(); diff --git a/src/common/function/src/scalars/vector/impl_conv.rs b/src/common/function/src/scalars/vector/impl_conv.rs index 219361a302c..00d23cedf7f 100644 --- a/src/common/function/src/scalars/vector/impl_conv.rs +++ b/src/common/function/src/scalars/vector/impl_conv.rs @@ -53,8 +53,10 @@ pub fn binlit_as_veclit(bytes: &[u8]) -> Result> { }) } else { let v = bytes - .chunks_exact(std::mem::size_of::()) - .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) + .as_chunks::<{ std::mem::size_of::() }>() + .0 + .iter() + .map(|chunk| f32::from_le_bytes(*chunk)) .collect::>(); Ok(Cow::Owned(v)) } diff --git a/src/common/grpc-expr/Cargo.toml b/src/common/grpc-expr/Cargo.toml index a689de7c018..b735e8f9ba8 100644 --- a/src/common/grpc-expr/Cargo.toml +++ b/src/common/grpc-expr/Cargo.toml @@ -9,7 +9,6 @@ workspace = true [dependencies] api.workspace = true -common-base.workspace = true common-catalog.workspace = true common-error.workspace = true common-macro.workspace = true @@ -21,6 +20,3 @@ prost.workspace = true snafu.workspace = true store-api.workspace = true table.workspace = true - -[dev-dependencies] -paste.workspace = true diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index d5214ae5ddb..ffa148c7d0b 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -24,21 +24,23 @@ use api::v1::{ SkippingIndexType as PbSkippingIndexType, column_def, }; use common_query::AddColumnLocation; +use datatypes::json::{JsonSettings, JsonTypeHint}; +use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use table::metadata::{TableId, TableMeta}; use table::requests::{ AddColumnRequest, AlterKind, AlterTableRequest, AnnotationFamily, ModifyColumnTypeRequest, - SetDefaultRequest, SetIndexOption, UnsetIndexOption, + SetDefaultRequest, SetIndexOption, SetJsonSettingsRequest, UnsetIndexOption, }; use crate::error::{ self, ColumnNotFoundSnafu, InvalidColumnDefSnafu, InvalidIndexOptionSnafu, - InvalidSetFulltextOptionRequestSnafu, InvalidSetSkippingIndexOptionRequestSnafu, - InvalidSetTableOptionRequestSnafu, InvalidUnsetTableOptionRequestSnafu, - MissingAlterIndexOptionSnafu, MissingFieldSnafu, MissingTableMetaSnafu, - MissingTimestampColumnSnafu, Result, UnknownLocationTypeSnafu, + InvalidJsonSettingsSnafu, InvalidSetFulltextOptionRequestSnafu, + InvalidSetSkippingIndexOptionRequestSnafu, InvalidSetTableOptionRequestSnafu, + InvalidUnsetTableOptionRequestSnafu, MissingAlterIndexOptionSnafu, MissingFieldSnafu, + MissingTableMetaSnafu, MissingTimestampColumnSnafu, Result, UnknownLocationTypeSnafu, }; const LOCATION_TYPE_FIRST: i32 = LocationType::First as i32; @@ -59,6 +61,34 @@ fn annotation_family_of_keys<'a>( }) } +fn json_settings_from_proto(settings: api::v1::JsonSettings) -> Result { + let type_hints = settings + .type_hints + .into_iter() + .map(|hint| { + let data_type = ConcreteDataType::from( + ColumnDataTypeWrapper::try_new(hint.data_type, hint.datatype_extension) + .context(error::ColumnDataTypeSnafu)?, + ); + + Ok(JsonTypeHint { + path: hint.path, + data_type, + // Index configuration is not supported yet, so this is temporarily + // hardcoded to false. + inverted_index: false, + }) + }) + .collect::>>()?; + + JsonSettings::try_new(type_hints, settings.max_auto_expanded_paths).map_err(|err| { + InvalidJsonSettingsSnafu { + err: err.to_string(), + } + .build() + }) +} + /// Returns the annotation family when `kind` is a SET/UNSET whose keys all /// belong to one family — the alters that only rewrite table metadata and skip /// region dispatch. A mixed batch is an error; never interpret it as "not an @@ -189,6 +219,17 @@ pub fn alter_expr_to_request( columns: modify_column_type_requests, } } + Kind::SetJsonSettings(set_json_settings) => { + let settings = set_json_settings + .settings + .context(MissingFieldSnafu { field: "settings" })?; + AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: set_json_settings.column_name, + settings: json_settings_from_proto(settings)?, + }, + } + } Kind::DropColumns(DropColumns { drop_columns }) => AlterKind::DropColumns { names: drop_columns.into_iter().map(|c| c.name).collect(), }, diff --git a/src/common/grpc-expr/src/error.rs b/src/common/grpc-expr/src/error.rs index ccc40efb5cb..6c9302c39f4 100644 --- a/src/common/grpc-expr/src/error.rs +++ b/src/common/grpc-expr/src/error.rs @@ -73,6 +73,13 @@ pub enum Error { source: api::error::Error, }, + #[snafu(display("Invalid JSON settings: {}", err))] + InvalidJsonSettings { + err: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unknown location type: {}", location_type))] UnknownLocationType { location_type: i32, @@ -188,6 +195,7 @@ impl ErrorExt for Error { | Error::MissingTimestampColumn { .. } => StatusCode::InvalidArguments, Error::MissingField { .. } => StatusCode::InvalidArguments, Error::InvalidColumnDef { source, .. } => source.status_code(), + Error::InvalidJsonSettings { .. } => StatusCode::InvalidArguments, Error::UnknownLocationType { .. } => StatusCode::InvalidArguments, Error::UnknownColumnDataType { .. } | Error::InvalidFulltextIndexColumnType { .. } => { diff --git a/src/common/mem-prof/Cargo.toml b/src/common/mem-prof/Cargo.toml index 10b67cf7fa7..46d959b02cc 100644 --- a/src/common/mem-prof/Cargo.toml +++ b/src/common/mem-prof/Cargo.toml @@ -26,3 +26,9 @@ 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.7" + +[package.metadata.cargo-shear] +# `tikv-jemalloc-sys` is a link-only dependency: it enables jemalloc +# profiling/stats features and pins the linked jemalloc instance, but is +# never referenced by name in source code. +ignored = ["tikv-jemalloc-sys"] diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 8deb3d00814..bd974deb4cf 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -46,7 +46,6 @@ common-time.workspace = true common-wal.workspace = true common-workload.workspace = true datafusion-common.workspace = true -datafusion-expr.workspace = true datatypes.workspace = true deadpool = { workspace = true, optional = true } deadpool-postgres = { workspace = true, optional = true } @@ -55,7 +54,6 @@ etcd-client.workspace = true flexbuffers = "25.2" futures.workspace = true futures-util.workspace = true -greptime-proto.workspace = true hex.workspace = true humantime-serde.workspace = true itertools.workspace = true diff --git a/src/common/meta/src/ddl/alter_table/executor.rs b/src/common/meta/src/ddl/alter_table/executor.rs index 53bc4b3d05b..419aa9ae12d 100644 --- a/src/common/meta/src/ddl/alter_table/executor.rs +++ b/src/common/meta/src/ddl/alter_table/executor.rs @@ -333,6 +333,7 @@ fn build_new_table_info( } AlterKind::DropColumns { .. } | AlterKind::ModifyColumnTypes { .. } + | AlterKind::SetJsonSettings { .. } | AlterKind::SetTableOptions { .. } | AlterKind::UnsetTableOptions { .. } | AlterKind::SetAnnotations { .. } diff --git a/src/common/meta/src/ddl/alter_table/region_request.rs b/src/common/meta/src/ddl/alter_table/region_request.rs index 73edd4c3615..b4f7b679422 100644 --- a/src/common/meta/src/ddl/alter_table/region_request.rs +++ b/src/common/meta/src/ddl/alter_table/region_request.rs @@ -90,6 +90,7 @@ fn create_proto_alter_kind( }))) } Kind::ModifyColumnTypes(x) => Ok(Some(alter_request::Kind::ModifyColumnTypes(x.clone()))), + Kind::SetJsonSettings(x) => Ok(Some(alter_request::Kind::SetJsonSettings(x.clone()))), Kind::DropColumns(x) => { let drop_columns = x .drop_columns diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index e36c89fbb34..89ade1b5cbe 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -42,7 +42,7 @@ use crate::cache_invalidator::Context; use crate::ddl::DdlContext; use crate::ddl::event::flow::{CREATE_FLOW_EVENT_TYPE, CreateFlowEventIntent, FlowDdlEvent}; use crate::ddl::utils::{add_peer_context_if_needed, map_to_procedure_error}; -use crate::error::{self, Result, UnexpectedSnafu}; +use crate::error::{self, InvalidFlowOptionSnafu, Result, UnexpectedSnafu}; use crate::instruction::{CacheIdent, CreateFlow, DropFlow}; use crate::key::flow::flow_info::{FlowInfoValue, FlowScheduleConfig, FlowStatus}; use crate::key::flow::flow_route::FlowRouteValue; @@ -558,9 +558,9 @@ pub fn validate_flow_options(flow_task: &CreateFlowTask) -> Result<()> { | FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY | FlowType::FLOW_TYPE_KEY => {} unknown => { - return UnexpectedSnafu { + return InvalidFlowOptionSnafu { err_msg: format!( - "Unknown flow option '{unknown}', supported user options: {DEFER_ON_MISSING_SOURCE_KEY}, {FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY}" + "unknown flow option '{unknown}', supported options: {DEFER_ON_MISSING_SOURCE_KEY}, {FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY}" ), } .fail(); diff --git a/src/common/meta/src/ddl/create_table/executor.rs b/src/common/meta/src/ddl/create_table/executor.rs index 88d88fe277c..332c9e57eb4 100644 --- a/src/common/meta/src/ddl/create_table/executor.rs +++ b/src/common/meta/src/ddl/create_table/executor.rs @@ -182,7 +182,19 @@ impl CreateTableExecutor { region_wal_options: RegionWalOptions, ) -> Result<()> { if !column_metadatas.is_empty() { - update_table_info_column_ids(&mut table_info, column_metadatas); + // Region metadata can include engine-internal columns that are not visible in + // the table schema. Only visible columns belong in the table's column IDs. + let visible_column_metadatas = column_metadatas + .iter() + .filter(|column| { + table_info + .meta + .schema + .contains_column(&column.column_schema.name) + }) + .cloned() + .collect::>(); + update_table_info_column_ids(&mut table_info, &visible_column_metadatas); } let detecting_regions = convert_region_routes_to_detecting_regions(&table_route.region_routes); diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs index 670f3c7c1dc..c1ab0e5e5a2 100644 --- a/src/common/meta/src/ddl/event/table.rs +++ b/src/common/meta/src/ddl/event/table.rs @@ -193,6 +193,7 @@ pub(crate) fn alter_table_kind_name(kind: &AlterTableKind) -> Option<&'static st AlterTableKind::DropColumns(_) => Some("drop_columns"), AlterTableKind::RenameTable(_) => Some("rename_table"), AlterTableKind::ModifyColumnTypes(_) => Some("modify_column_types"), + AlterTableKind::SetJsonSettings(_) => Some("set_json_settings"), AlterTableKind::SetTableOptions(_) => Some("set_table_options"), AlterTableKind::UnsetTableOptions(_) => Some("unset_table_options"), AlterTableKind::SetIndex(_) => Some("set_index"), diff --git a/src/common/meta/src/ddl/tests/create_flow.rs b/src/common/meta/src/ddl/tests/create_flow.rs index e7004656fe9..ad334e29863 100644 --- a/src/common/meta/src/ddl/tests/create_flow.rs +++ b/src/common/meta/src/ddl/tests/create_flow.rs @@ -18,6 +18,8 @@ use std::sync::Arc; use api::v1::flow::CreateRequest; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; use common_procedure::{Procedure, Status}; use common_procedure_test::execute_procedure_until_done; use table::table_name::TableName; @@ -326,6 +328,26 @@ fn test_validate_flow_options_allows_incremental_read_option() { validate_flow_options(&task).unwrap(); } +#[test] +fn test_validate_flow_options_unknown_option_matches_invalid_sql_error() { + let mut task = test_create_flow_task( + "my_flow", + vec![], + TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "my_sink_table"), + false, + ); + task.flow_options + .insert("foo".to_string(), "sensitive-value".to_string()); + + let err = validate_flow_options(&task).unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert_eq!( + "Invalid SQL, error: unknown flow option 'foo', supported options: defer_on_missing_source, experimental_enable_incremental_read", + err.output_msg() + ); + assert!(!err.output_msg().contains("sensitive-value")); +} + #[test] fn test_validate_flow_options_rejects_schedule_and_internal_keys_as_unknown() { for key in [ @@ -346,9 +368,10 @@ fn test_validate_flow_options_rejects_schedule_and_internal_keys_as_unknown() { .insert(key.to_string(), "value".to_string()); let err = validate_flow_options(&task).unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); assert!( err.to_string() - .contains(&format!("Unknown flow option '{key}'")), + .contains(&format!("unknown flow option '{key}'")), "unexpected error for {key}: {err}" ); } @@ -884,10 +907,10 @@ async fn test_create_flow_rejects_unknown_option_in_meta_task() { let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); - assert_matches!(err, error::Error::Unexpected { .. }); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); assert!( err.to_string() - .contains("Unknown flow option 'unknown_option'") + .contains("unknown flow option 'unknown_option'") ); } diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3f4868ffba0..5c1febb9b4d 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -355,13 +355,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Corrupted table route data, err: {}", err_msg))] - RouteInfoCorrupted { - err_msg: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Illegal state from server, code: {}, error: {}", code, err_msg))] IllegalServerState { code: i32, @@ -522,6 +515,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid SQL, error: {err_msg}"))] + InvalidFlowOption { + err_msg: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid flow request body: {:?}", body))] InvalidFlowRequestBody { body: Box>, @@ -1213,7 +1213,6 @@ impl ErrorExt for Error { SerdeJson { .. } | ParseOption { .. } - | RouteInfoCorrupted { .. } | InvalidProtoMsg { .. } | InvalidMetadata { .. } | Unexpected { .. } @@ -1255,6 +1254,7 @@ impl ErrorExt for Error { ProcedureNotFound { .. } | InvalidViewInfo { .. } + | InvalidFlowOption { .. } | PrimaryKeyNotFound { .. } | EmptyKey { .. } | AlterLogicalTablesInvalidArguments { .. } diff --git a/src/common/meta/src/instruction.rs b/src/common/meta/src/instruction.rs index 60ebe14e37c..6ae074b0ac9 100644 --- a/src/common/meta/src/instruction.rs +++ b/src/common/meta/src/instruction.rs @@ -732,12 +732,12 @@ impl PackedRegionFileRefs { } let mut refs = HashSet::new(); - for bytes in files.chunks_exact(16) { + for bytes in files.as_chunks::<16>().0 { 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) { + for bytes in indexed.as_chunks::<24>().0 { let mut id = [0; 16]; id.copy_from_slice(&bytes[..16]); let mut version = [0; 8]; diff --git a/src/common/meta/src/key/schema_name.rs b/src/common/meta/src/key/schema_name.rs index 689f72ddb32..dcff9c6b338 100644 --- a/src/common/meta/src/key/schema_name.rs +++ b/src/common/meta/src/key/schema_name.rs @@ -334,6 +334,21 @@ impl SchemaManager { Box::pin(stream) } + + /// Returns schema names and values belonging to the target `catalog`. + /// Legacy `null` values are returned as [`SchemaNameValue::default()`]. + pub fn schemas(&self, catalog: &str) -> BoxStream<'static, Result<(String, SchemaNameValue)>> { + let start_key = SchemaNameKey::range_start_key(catalog); + let req = RangeRequest::new().with_prefix(start_key.as_bytes()); + + let stream = PaginationStream::new(self.kv_backend.clone(), req, DEFAULT_PAGE_SIZE, |kv| { + let value = SchemaNameValue::try_from_raw_value(&kv.value)?.unwrap_or_default(); + Ok((schema_decoder(kv)?, value)) + }) + .into_stream(); + + Box::pin(stream) + } } #[derive(Debug, Clone, Hash, Eq, PartialEq, Deserialize, Serialize)] @@ -358,12 +373,124 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; + use futures::TryStreamExt; use store_api::mito_engine_options::{ TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, TWCS_TRIGGER_FILE_NUM, }; use super::*; + use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; + use crate::rpc::store::PutRequest; + + #[tokio::test] + async fn test_schemas() { + let manager = SchemaManager::new(Arc::new(MemoryKvBackend::default())); + assert!( + manager + .schemas("catalog") + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + + let mut expected = BTreeMap::new(); + for i in 0..=DEFAULT_PAGE_SIZE { + let name = format!("schema_{i}"); + let value = if i == 0 { + SchemaNameValue::default() + } else { + SchemaNameValue { + ttl: Some(Duration::from_secs(i as u64).into()), + extra_options: BTreeMap::from([("foo".to_string(), i.to_string())]), + create_procedure_id: Some(format!("procedure_{i}")), + } + }; + manager + .create( + SchemaNameKey::new("catalog", &name), + Some(value.clone()), + false, + ) + .await + .unwrap(); + expected.insert(name, value); + } + manager + .create(SchemaNameKey::new("catalog_other", "schema"), None, false) + .await + .unwrap(); + + let schemas = manager + .schemas("catalog") + .try_collect::>() + .await + .unwrap(); + assert_eq!(schemas, expected.into_iter().collect::>()); + assert_eq!( + manager + .schema_names("catalog") + .try_collect::>() + .await + .unwrap(), + schemas + .into_iter() + .map(|(name, _)| name) + .collect::>() + ); + } + + #[tokio::test] + async fn test_schemas_legacy_and_invalid_values() { + let kv_backend = Arc::new(MemoryKvBackend::default()); + let manager = SchemaManager::new(kv_backend.clone()); + let key = SchemaNameKey::new("catalog", "schema").to_bytes(); + + for (raw, expected) in [ + (b"null".as_slice(), SchemaNameValue::default()), + ( + br#"{"ttl":"10s"}"#.as_slice(), + SchemaNameValue { + ttl: Some(Duration::from_secs(10).into()), + ..Default::default() + }, + ), + ] { + kv_backend + .put(PutRequest::new().with_key(key.clone()).with_value(raw)) + .await + .unwrap(); + assert_eq!( + manager + .schemas("catalog") + .try_collect::>() + .await + .unwrap(), + vec![("schema".to_string(), expected)] + ); + } + + kv_backend + .put(PutRequest::new().with_key(key).with_value(b"invalid")) + .await + .unwrap(); + assert!( + manager + .schemas("catalog") + .try_collect::>() + .await + .is_err() + ); + assert_eq!( + manager + .schema_names("catalog") + .try_collect::>() + .await + .unwrap(), + vec!["schema".to_string()] + ); + } #[test] fn test_display_schema_value() { diff --git a/src/common/meta/src/kv_backend/chroot.rs b/src/common/meta/src/kv_backend/chroot.rs index 1a7e7c615ac..4d2d16d0f2f 100644 --- a/src/common/meta/src/kv_backend/chroot.rs +++ b/src/common/meta/src/kv_backend/chroot.rs @@ -82,7 +82,7 @@ impl KvBackend for ChrootKvBackend { async fn batch_put(&self, mut req: BatchPutRequest) -> Result { for kv in req.kvs.iter_mut() { - kv.key = self.key_prepend_root(kv.key.drain(..).collect()); + kv.key = self.key_prepend_root(std::mem::take(&mut kv.key)); } let mut res = self.inner.batch_put(req).await?; res.prev_kvs = res diff --git a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs index 27950f0fd5d..8defc6bd0f2 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs @@ -25,8 +25,8 @@ use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; use crate::reconciliation::reconcile_table::update_table_info::UpdateTableInfo; use crate::reconciliation::reconcile_table::{ReconcileTableContext, State, TableMetadataState}; use crate::reconciliation::utils::{ - ResolveColumnMetadataResult, build_column_metadata_from_table_info, - check_column_metadatas_consistent, resolve_column_metadatas_with_latest, + ResolveColumnMetadataResult, build_reconciliation_column_metadata, + check_column_metadatas_consistent, reorder_tag_columns, resolve_column_metadatas_with_latest, resolve_column_metadatas_with_metasrv, }; @@ -92,6 +92,8 @@ impl State for ResolveColumnMetadata { ctx.persistent_ctx.table_info_value = Some(table_info_value); if let Some(column_metadatas) = check_column_metadatas_consistent(&self.region_metadata) { + let column_metadatas = + reorder_tag_columns(&column_metadatas, &self.region_metadata[0].primary_key)?; // Safety: fetched in the above. let table_info_value = ctx.persistent_ctx.table_info_value.clone().unwrap(); info!( @@ -125,7 +127,7 @@ impl State for ResolveColumnMetadata { .table_info .name_to_ids() .context(MissingColumnIdsSnafu)?; - let column_metadata = build_column_metadata_from_table_info( + let column_metadata = build_reconciliation_column_metadata( table_info_value.table_info.meta.schema.column_schemas(), &table_info_value.table_info.meta.primary_key_indices, &name_to_ids, diff --git a/src/common/meta/src/reconciliation/utils.rs b/src/common/meta/src/reconciliation/utils.rs index 6ddc084596b..326516ecd80 100644 --- a/src/common/meta/src/reconciliation/utils.rs +++ b/src/common/meta/src/reconciliation/utils.rs @@ -85,6 +85,70 @@ pub(crate) fn check_column_metadatas_consistent( Some(region_metadatas[0].column_metadatas.clone()) } +/// Returns columns with tag values in primary-key order while preserving the +/// positions of non-tag columns. +pub(crate) fn reorder_tag_columns( + column_metadatas: &[ColumnMetadata], + primary_key: &[u32], +) -> Result> { + let tag_count = column_metadatas + .iter() + .filter(|column| column.semantic_type == SemanticType::Tag) + .count(); + ensure!( + primary_key.len() == tag_count, + UnexpectedSnafu { + err_msg: format!( + "Number of primary key columns {} does not match tag columns {}", + primary_key.len(), + tag_count, + ), + } + ); + + let mut primary_key_ids = HashSet::with_capacity(primary_key.len()); + let mut tags = Vec::with_capacity(primary_key.len()); + for column_id in primary_key { + let column = column_metadatas + .iter() + .find(|column| column.column_id == *column_id) + .with_context(|| UnexpectedSnafu { + err_msg: format!( + "Primary key column {} not found in column metadata", + column_id + ), + })?; + ensure!( + column.semantic_type == SemanticType::Tag, + UnexpectedSnafu { + err_msg: format!("Primary key column {} is not a tag", column_id), + } + ); + ensure!( + primary_key_ids.insert(column_id), + UnexpectedSnafu { + err_msg: format!("Primary key column {} is duplicated", column_id), + } + ); + tags.push(column); + } + + let mut tags = tags.into_iter(); + + column_metadatas + .iter() + .map(|column| { + if column.semantic_type == SemanticType::Tag { + tags.next().cloned().with_context(|| UnexpectedSnafu { + err_msg: "Primary key has fewer columns than tags".to_string(), + }) + } else { + Ok(column.clone()) + } + }) + .collect() +} + /// Resolves column metadata inconsistencies among the given region metadatas /// by using the column metadata from the metasrv as the source of truth. /// @@ -157,8 +221,13 @@ pub(crate) fn resolve_column_metadatas_with_latest( } } - // TODO(weny): verify the new column metadatas are acceptable for regions. - Ok((latest_region_metadata.column_metadatas.clone(), region_ids)) + Ok(( + reorder_tag_columns( + &latest_region_metadata.column_metadatas, + &latest_region_metadata.primary_key, + )?, + region_ids, + )) } /// Constructs a vector of [`ColumnMetadata`] from the provided table information. @@ -206,6 +275,22 @@ pub(crate) fn build_column_metadata_from_table_info( .collect::>>() } +/// Builds SyncColumns metadata with tags in the TableInfo primary-key order. +/// Non-tag columns retain their schema positions. +pub(crate) fn build_reconciliation_column_metadata( + column_schemas: &[ColumnSchema], + primary_key_indexes: &[usize], + name_to_ids: &HashMap, +) -> Result> { + let column_metadatas = + build_column_metadata_from_table_info(column_schemas, primary_key_indexes, name_to_ids)?; + let primary_key = primary_key_indexes + .iter() + .map(|index| column_metadatas[*index].column_id) + .collect::>(); + reorder_tag_columns(&column_metadatas, &primary_key) +} + /// Checks whether the schema invariants hold between the existing and new column metadata. /// /// Invariants: @@ -1180,6 +1265,104 @@ mod tests { assert_matches!(err, Error::MissingColumnInColumnMetadata { .. }); } + #[test] + fn test_reconciliation_columns_follow_primary_key_order() { + let columns = vec![ + ColumnMetadata { + column_schema: ColumnSchema::new( + "tag_b", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 1, + }, + ColumnMetadata { + column_schema: ColumnSchema::new("field", ConcreteDataType::int32_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 2, + }, + ColumnMetadata { + column_schema: ColumnSchema::new( + "tag_a", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 3, + }, + ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + semantic_type: SemanticType::Timestamp, + column_id: 4, + }, + ]; + let schemas = columns + .iter() + .map(|column| column.column_schema.clone()) + .collect::>(); + let ids = columns + .iter() + .map(|column| (column.column_schema.name.clone(), column.column_id)) + .collect(); + let expected = vec![3, 2, 1, 4]; + assert_eq!( + build_reconciliation_column_metadata(&schemas, &[2, 0], &ids) + .unwrap() + .iter() + .map(|column| column.column_id) + .collect::>(), + expected + ); + + let mut metadata = build_region_metadata(RegionId::new(1024, 0), &columns); + metadata.primary_key = vec![3, 1]; + metadata.schema_version = 2; + + assert_eq!( + check_column_metadatas_consistent(&[metadata.clone()]).unwrap(), + columns + ); + + assert_eq!( + resolve_column_metadatas_with_latest(&[metadata]) + .unwrap() + .0 + .iter() + .map(|column| column.column_id) + .collect::>(), + expected + ); + } + + #[test] + fn test_reorder_tag_columns_rejects_invalid_primary_key() { + let columns = new_test_column_metadatas(); + + let err = reorder_tag_columns(&columns, &[999]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("not found in column metadata")); + + let err = reorder_tag_columns(&columns, &[]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("does not match tag columns")); + + let err = reorder_tag_columns(&columns, &[2]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("is not a tag")); + + let mut duplicate_tags = columns.clone(); + duplicate_tags[2].semantic_type = SemanticType::Tag; + let err = reorder_tag_columns(&duplicate_tags, &[0, 0]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("is duplicated")); + } + #[test] fn test_check_column_metadatas_consistent() { let column_metadatas = new_test_column_metadatas(); diff --git a/src/common/meta/src/rpc/router.rs b/src/common/meta/src/rpc/router.rs index 519b12e081f..20fcd955abb 100644 --- a/src/common/meta/src/rpc/router.rs +++ b/src/common/meta/src/rpc/router.rs @@ -14,22 +14,15 @@ use std::collections::{BTreeMap, HashMap, HashSet}; -use api::v1::meta::{ - Partition as PbPartition, Peer as PbPeer, Region as PbRegion, Table as PbTable, - TableRoute as PbTableRoute, -}; use common_time::util::current_time_millis; use derive_builder::Builder; use serde::ser::SerializeSeq; use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use snafu::OptionExt; use store_api::region_engine::RegionRole; use store_api::storage::{RegionId, RegionNumber}; use strum::AsRefStr; -use table::table_name::TableName; use crate::DatanodeId; -use crate::error::{self, Result}; use crate::key::RegionDistribution; use crate::peer::Peer; @@ -55,20 +48,13 @@ pub fn region_distribution(region_routes: &[RegionRoute]) -> RegionDistribution .add_follower_region(region_number); } } - for (_, region_role_set) in regions_id_map.iter_mut() { + for region_role_set in regions_id_map.values_mut() { // Sort the regions in ascending order. region_role_set.sort() } regions_id_map } -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] -pub struct TableRoute { - pub table: Table, - pub region_routes: Vec, - region_leaders: HashMap>, -} - /// Returns the leader peers of the table. pub fn find_leaders(region_routes: &[RegionRoute]) -> HashSet { region_routes @@ -167,94 +153,6 @@ pub fn find_follower_regions(region_routes: &[RegionRoute], datanode: &Peer) -> .collect() } -impl TableRoute { - pub fn new(table: Table, region_routes: Vec) -> Self { - let region_leaders = region_routes - .iter() - .map(|x| (x.region.id.region_number(), x.leader_peer.clone())) - .collect::>(); - Self { - table, - region_routes, - region_leaders, - } - } - - pub fn try_from_raw(peers: &[PbPeer], table_route: PbTableRoute) -> Result { - let table = table_route - .table - .context(error::RouteInfoCorruptedSnafu { - err_msg: "'table' is empty in table route", - })? - .try_into()?; - - let mut region_routes = Vec::with_capacity(table_route.region_routes.len()); - for region_route in table_route.region_routes.into_iter() { - let region = region_route - .region - .context(error::RouteInfoCorruptedSnafu { - err_msg: "'region' is empty in region route", - })? - .into(); - - let leader_peer = peers.get(region_route.leader_peer_index as usize).cloned(); - - let follower_peers = region_route - .follower_peer_indexes - .into_iter() - .filter_map(|x| peers.get(x as usize).cloned()) - .collect::>(); - - region_routes.push(RegionRoute { - region, - leader_peer, - follower_peers, - leader_state: None, - leader_down_since: None, - write_route_policy: None, - }); - } - - Ok(Self::new(table, region_routes)) - } -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] -pub struct Table { - pub id: u64, - pub table_name: TableName, - #[serde(serialize_with = "as_utf8", deserialize_with = "from_utf8")] - pub table_schema: Vec, -} - -impl TryFrom for Table { - type Error = error::Error; - - fn try_from(t: PbTable) -> Result { - let table_name = t - .table_name - .context(error::RouteInfoCorruptedSnafu { - err_msg: "table name required", - })? - .into(); - Ok(Self { - id: t.id, - table_name, - table_schema: t.table_schema, - }) - } -} - -impl From for PbTable { - fn from(table: Table) -> Self { - PbTable { - id: table.id, - table_name: Some(table.table_name.into()), - table_schema: table.table_schema, - } - } -} - #[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Builder)] pub struct RegionRoute { pub region: Region, @@ -430,18 +328,6 @@ impl RegionRoute { } } -pub struct RegionRoutes(pub Vec); - -impl RegionRoutes { - pub fn region_leader_map(&self) -> HashMap { - convert_to_region_leader_map(&self.0) - } - - pub fn find_region_leader(&self, region_number: RegionNumber) -> Option<&Peer> { - self.region_leader_map().get(®ion_number).copied() - } -} - #[derive(Debug, Clone, Default, PartialEq, Serialize)] pub struct Region { pub id: RegionId, @@ -508,49 +394,6 @@ impl Region { } } -/// Gets the partition expression of the `PbRegion` in compatible mode. -#[allow(deprecated)] -pub fn pb_region_partition_expr(r: &PbRegion) -> String { - if let Some(partition) = &r.partition { - if !partition.expression.is_empty() { - partition.expression.clone() - } else if !partition.value_list.is_empty() { - String::from_utf8_lossy(&partition.value_list[0]).to_string() - } else { - "".to_string() - } - } else { - "".to_string() - } -} - -impl From for Region { - fn from(r: PbRegion) -> Self { - let partition_expr = pb_region_partition_expr(&r); - Self { - id: r.id.into(), - name: r.name, - partition_expr, - attrs: r.attrs.into_iter().collect::>(), - } - } -} - -impl From for PbRegion { - fn from(region: Region) -> Self { - let partition_expr = region.partition_expr(); - Self { - id: region.id.into(), - name: region.name, - partition: Some(PbPartition { - expression: partition_expr, - ..Default::default() - }), - attrs: region.attrs.into_iter().collect::>(), - } - } -} - #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] pub struct LegacyPartition { #[serde(serialize_with = "as_utf8_vec", deserialize_with = "from_utf8_vec")] @@ -559,23 +402,6 @@ pub struct LegacyPartition { pub value_list: Vec>, } -fn as_utf8(val: &[u8], serializer: S) -> std::result::Result { - serializer.serialize_str( - String::from_utf8(val.to_vec()) - .unwrap_or_else(|_| "".to_string()) - .as_str(), - ) -} - -pub fn from_utf8<'de, D>(deserializer: D) -> std::result::Result, D::Error> -where - D: Deserializer<'de>, -{ - let s = String::deserialize(deserializer)?; - - Ok(s.into_bytes()) -} - fn as_utf8_vec( val: &[Vec], serializer: S, @@ -888,58 +714,4 @@ mod tests { assert_eq!(got, p); } - - #[test] - #[allow(deprecated)] - fn test_region_partition_expr() { - let r = PbRegion { - id: 1, - name: "r1".to_string(), - partition: None, - attrs: Default::default(), - }; - assert_eq!(pb_region_partition_expr(&r), ""); - - let r2: Region = r.into(); - assert_eq!(r2.partition_expr(), ""); - - let r3: PbRegion = r2.into(); - assert_eq!(r3.partition.as_ref().unwrap().expression, ""); - - let r = PbRegion { - id: 1, - name: "r1".to_string(), - partition: Some(PbPartition { - column_list: vec![b"a".to_vec()], - value_list: vec![b"{}".to_vec()], - expression: Default::default(), - }), - attrs: Default::default(), - }; - assert_eq!(pb_region_partition_expr(&r), "{}"); - - let r2: Region = r.into(); - assert_eq!(r2.partition_expr(), "{}"); - - let r3: PbRegion = r2.into(); - assert_eq!(r3.partition.as_ref().unwrap().expression, "{}"); - - let r = PbRegion { - id: 1, - name: "r1".to_string(), - partition: Some(PbPartition { - column_list: vec![b"a".to_vec()], - value_list: vec![b"{}".to_vec()], - expression: "a>b".to_string(), - }), - attrs: Default::default(), - }; - assert_eq!(pb_region_partition_expr(&r), "a>b"); - - let r2: Region = r.into(); - assert_eq!(r2.partition_expr(), "a>b"); - - let r3: PbRegion = r2.into(); - assert_eq!(r3.partition.as_ref().unwrap().expression, "a>b"); - } } diff --git a/src/common/meta/src/rpc/store.rs b/src/common/meta/src/rpc/store.rs index 25536993972..a9c36199761 100644 --- a/src/common/meta/src/rpc/store.rs +++ b/src/common/meta/src/rpc/store.rs @@ -199,7 +199,7 @@ impl RangeResponse { } pub fn take_kvs(&mut self) -> Vec { - self.kvs.drain(..).collect() + std::mem::take(&mut self.kvs) } } @@ -451,7 +451,7 @@ impl BatchPutResponse { } pub fn take_prev_kvs(&mut self) -> Vec { - self.prev_kvs.drain(..).collect() + std::mem::take(&mut self.prev_kvs) } } @@ -780,7 +780,7 @@ impl DeleteRangeResponse { } pub fn take_prev_kvs(&mut self) -> Vec { - self.prev_kvs.drain(..).collect() + std::mem::take(&mut self.prev_kvs) } } diff --git a/src/common/meta/src/sequence.rs b/src/common/meta/src/sequence.rs index cd0a8ebf881..9d56170833c 100644 --- a/src/common/meta/src/sequence.rs +++ b/src/common/meta/src/sequence.rs @@ -188,7 +188,7 @@ impl Inner { } error::NextSequenceSnafu { - err_msg: format!("{}.next()", &self.name), + err_msg: format!("{}.next()", self.name), } .fail() } @@ -264,7 +264,7 @@ impl Inner { } error::NextSequenceSnafu { - err_msg: format!("{}.next_range()", &self.name), + err_msg: format!("{}.next_range()", self.name), } .fail() } diff --git a/src/common/meta/src/snapshot.rs b/src/common/meta/src/snapshot.rs index f35c86e2e7b..3582b0acecb 100644 --- a/src/common/meta/src/snapshot.rs +++ b/src/common/meta/src/snapshot.rs @@ -408,7 +408,6 @@ mod tests { async fn test_dump_and_restore() { common_telemetry::init_default_ut_logging(); let (temp_dir, kv_backend, manager) = test_env("test_dump_and_restore"); - let temp_path = temp_dir.path(); for i in 0..10 { kv_backend @@ -420,32 +419,16 @@ mod tests { .await .unwrap(); } - let dump_path = temp_path.join("snapshot"); - manager - .dump(&format!( - "{}/metadata_snapshot", - &dump_path.as_path().display().to_string() - )) - .await - .unwrap(); + manager.dump("snapshot/metadata_snapshot").await.unwrap(); + let snapshot_key = "snapshot/metadata_snapshot.metadata.fb"; + assert!(temp_dir.path().join("data").join(snapshot_key).is_file()); // Clean up the kv backend kv_backend.clear(); - let err = manager - .dump(&format!( - "{}/metadata_snapshot.metadata.fb", - &dump_path.as_path().display().to_string() - )) - .await - .unwrap_err(); + let err = manager.dump(snapshot_key).await.unwrap_err(); assert_matches!(err, Error::Unexpected { .. }); assert!(err.to_string().contains("already exists")); - let restore_path = dump_path - .join("metadata_snapshot.metadata.fb") - .as_path() - .display() - .to_string(); - manager.restore(&restore_path).await.unwrap(); + manager.restore(snapshot_key).await.unwrap(); for i in 0..10 { let key = format!("test_{}", i); @@ -456,14 +439,11 @@ mod tests { #[tokio::test] async fn test_restore_from_nonexistent_file() { - let (temp_dir, _kv_backend, manager) = test_env("test_restore_from_nonexistent_file"); - let restore_path = temp_dir - .path() - .join("nonexistent.metadata.fb") - .as_path() - .display() - .to_string(); - let err = manager.restore(&restore_path).await.unwrap_err(); + let (_temp_dir, _kv_backend, manager) = test_env("test_restore_from_nonexistent_file"); + let err = manager + .restore("nonexistent.metadata.fb") + .await + .unwrap_err(); assert_matches!(err, Error::ReadObject { .. }) } } diff --git a/src/common/procedure-test/Cargo.toml b/src/common/procedure-test/Cargo.toml index e445a235bc9..907035dd052 100644 --- a/src/common/procedure-test/Cargo.toml +++ b/src/common/procedure-test/Cargo.toml @@ -10,5 +10,4 @@ workspace = true [dependencies] async-trait.workspace = true common-procedure = { workspace = true, features = ["testing"] } -snafu.workspace = true tokio.workspace = true diff --git a/src/common/query/Cargo.toml b/src/common/query/Cargo.toml index 9706098826f..4acdfafcb6c 100644 --- a/src/common/query/Cargo.toml +++ b/src/common/query/Cargo.toml @@ -39,3 +39,8 @@ common-base.workspace = true futures-util.workspace = true serde_json.workspace = true tokio.workspace = true + +[package.metadata.cargo-shear] +# `sqlparser` is required by `sqlparser_derive`: the generated code references +# `::sqlparser` paths, so it must stay a direct dependency. +ignored = ["sqlparser"] diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index 8ce028ef9fe..8151dfc5731 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -25,10 +25,13 @@ pub mod stream; pub mod test_util; use std::fmt::{Debug, Display, Formatter}; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use api::greptime_proto::v1::AddColumnLocation as Location; use api::greptime_proto::v1::add_column_location::LocationType; +use common_error::ext::BoxedError; use common_recordbatch::{ RecordBatches, SendableRecordBatchMapper, SendableRecordBatchStream, map_dictionary_to_values, map_dictionary_to_values_schema, @@ -74,6 +77,30 @@ pub struct OutputMeta { /// May exist for query output. One can retrieve execution metrics from this plan. pub plan: Option>, pub cost: OutputCost, + /// Process-local storage results attached by asynchronous ordinary-table batching. + /// Callers needing confirmed writes must await these. Dropping them does not cancel writes. + pub write_completions: Vec, +} + +/// Actual storage completion of an already accepted asynchronous write. +pub struct WriteCompletion(Pin> + Send>>); + +impl WriteCompletion { + /// Wraps the storage result without changing how the write is scheduled. + pub fn new(completion: impl Future> + Send + 'static) -> Self { + Self(Box::pin(completion)) + } + + /// Waits for storage success or failure, independently of queue admission. + pub async fn wait(self) -> Result<(), BoxedError> { + self.0.await + } +} + +impl Debug for WriteCompletion { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WriteCompletion").finish_non_exhaustive() + } } impl Output { @@ -163,18 +190,26 @@ impl Debug for OutputData { impl OutputMeta { pub fn new(plan: Option>, cost: usize) -> Self { - Self { plan, cost } + Self { + plan, + cost, + ..Default::default() + } } pub fn new_with_plan(plan: Arc) -> Self { Self { plan: Some(plan), cost: 0, + ..Default::default() } } pub fn new_with_cost(cost: usize) -> Self { - Self { plan: None, cost } + Self { + cost, + ..Default::default() + } } } diff --git a/src/common/query/src/stream.rs b/src/common/query/src/stream.rs index 0a46164d9b4..d6f1735e06c 100644 --- a/src/common/query/src/stream.rs +++ b/src/common/query/src/stream.rs @@ -95,7 +95,7 @@ impl DisplayAs for StreamScanAdapter { f, "StreamScanAdapter: [], schema: [" )?; - write!(f, "{:?}", &self.arrow_schema)?; + write!(f, "{:?}", self.arrow_schema)?; write!(f, "]") } } diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index a7f6e4a10ac..f8d645bd3e1 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -570,7 +570,7 @@ impl MetricCollector { } impl ExecutionPlanVisitor for MetricCollector { - type Error = !; + type Error = std::convert::Infallible; fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> std::result::Result { // skip if no metric available diff --git a/src/common/recordbatch/src/filter.rs b/src/common/recordbatch/src/filter.rs index 26c28aa3fa0..60ff399601f 100644 --- a/src/common/recordbatch/src/filter.rs +++ b/src/common/recordbatch/src/filter.rs @@ -12,215 +12,174 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Util record batch stream wrapper that can perform precise filter. +//! Precise single-column predicates and record batch filtering. + +mod predicate; +#[cfg(test)] +mod tests; +mod timestamp; use std::sync::Arc; -use common_time::timestamp::div_mod_units; use datafusion::error::Result as DfResult; use datafusion::logical_expr::{Expr, Literal, Operator}; use datafusion::physical_plan::PhysicalExpr; -use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar}; +use datafusion_common::arrow::array::{ArrayRef, Datum}; use datafusion_common::arrow::buffer::BooleanBuffer; -use datafusion_common::arrow::compute::kernels::cmp; -use datafusion_common::cast::{as_boolean_array, as_null_array, as_string_array}; +use datafusion_common::cast::{as_boolean_array, as_null_array}; use datafusion_common::{DataFusionError, ScalarValue, internal_err}; -use datatypes::arrow::array::{ - Array, ArrayAccessor, ArrayData, BooleanArray, BooleanBufferBuilder, DictionaryArray, - RecordBatch, StringArrayType, -}; +use datatypes::arrow::array::{Array, BooleanArray, BooleanBufferBuilder, RecordBatch}; use datatypes::arrow::compute::filter_record_batch; -use datatypes::arrow::datatypes::{DataType, TimeUnit, UInt32Type}; -use datatypes::arrow::error::ArrowError; -use datatypes::compute::or_kleene; +use datatypes::arrow::datatypes::DataType; use datatypes::data_type::{ConcreteDataType, DataType as _}; use datatypes::value::Value; use datatypes::vectors::VectorRef; -use regex::Regex; use snafu::ResultExt; -use crate::error::{ArrowComputeSnafu, Result, ToArrowScalarSnafu, UnsupportedOperationSnafu}; +pub use self::predicate::{regexp_is_match_dictionary, regexp_is_match_scalar}; +pub use self::timestamp::timestamp_scalar_value; +use crate::error::{Result, ToArrowScalarSnafu}; +use crate::filter::predicate::SimplePredicate; -/// An inplace expr evaluator for simple filter. Only support -/// - `col` `op` `literal` -/// - `literal` `op` `col` -/// -/// And the `op` is one of `=`, `!=`, `>`, `>=`, `<`, `<=`, -/// or regex operators: `~`, `~*`, `!~`, `!~*`. -/// -/// This struct contains normalized predicate expr. In the form of -/// `col` `op` `literal` where the `col` is provided from input. +/// Evaluates supported single-column predicates with precompiled regular expressions. +/// Compound predicates preserve SQL null semantics until conversion to a scan mask. #[derive(Debug, Clone)] pub struct SimpleFilterEvaluator { - /// Name of the referenced column. column_name: String, - /// The literal value. - literal: Scalar, - /// The operator. - op: Operator, - /// Only used when the operator is `Or`-chain. - literal_list: Vec>, - /// Pre-compiled regex. - /// Only used when the operator is regex operators. - /// If the regex is empty, it is also `None`. - regex: Option, - /// Whether the regex is negative. - regex_negative: bool, + predicate: SimplePredicate, } impl SimpleFilterEvaluator { pub fn new(column_name: String, lit: T, op: Operator) -> Option { - match op { + if !matches!( + op, Operator::Eq - | Operator::NotEq - | Operator::Lt - | Operator::LtEq - | Operator::Gt - | Operator::GtEq => {} - _ => return None, + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + ) { + return None; } - - let Expr::Literal(val, _) = lit.lit() else { + let Expr::Literal(value, _) = lit.lit() else { return None; }; - Some(Self { column_name, - literal: val.to_scalar().ok()?, - op, - literal_list: vec![], - regex: None, - regex_negative: false, + predicate: SimplePredicate::Comparison { + literal: value.to_scalar().ok()?, + op, + regex: None, + regex_negative: false, + }, }) } pub fn try_new(predicate: &Expr) -> Option { - match predicate { - Expr::BinaryExpr(binary) => { - // check if the expr is in the supported form - match binary.op { - Operator::Eq - | Operator::NotEq - | Operator::Lt - | Operator::LtEq - | Operator::Gt - | Operator::GtEq - | Operator::RegexMatch - | Operator::RegexIMatch - | Operator::RegexNotMatch - | Operator::RegexNotIMatch => {} - Operator::Or => { - let lhs = Self::try_new(&binary.left)?; - let rhs = Self::try_new(&binary.right)?; - if lhs.column_name != rhs.column_name - || !matches!(lhs.op, Operator::Eq | Operator::Or) - || !matches!(rhs.op, Operator::Eq | Operator::Or) - { - return None; - } - let mut list = vec![]; - let placeholder_literal = lhs.literal.clone(); - // above check guarantees the op is either `Eq` or `Or` - if matches!(lhs.op, Operator::Or) { - list.extend(lhs.literal_list); - } else { - list.push(lhs.literal); - } - if matches!(rhs.op, Operator::Or) { - list.extend(rhs.literal_list); - } else { - list.push(rhs.literal); - } - return Some(Self { - column_name: lhs.column_name, - literal: placeholder_literal, - op: Operator::Or, - literal_list: list, - regex: None, - regex_negative: false, - }); - } - _ => return None, - } - - // swap the expr if it is in the form of `literal` `op` `col` - let mut op = binary.op; - let (lhs, rhs) = match (&*binary.left, &*binary.right) { - (Expr::Column(col), Expr::Literal(lit, _)) => (col, lit), - (Expr::Literal(lit, _), Expr::Column(col)) => { - // safety: The previous check ensures the operator is able to swap. - op = op.swap().unwrap(); - (col, lit) - } - _ => return None, - }; - - let (regex, regex_negative) = Self::maybe_build_regex(op, rhs).ok()?; - let literal = rhs.to_scalar().ok()?; - Some(Self { - column_name: lhs.name.clone(), - literal, - op, - literal_list: vec![], - regex, - regex_negative, - }) - } - _ => None, - } + Self::try_new_with_column_type(predicate, &|_| None) } - /// Get the name of the referenced column. + /// Recognizes single-column filters, including lossless string dictionary casts + /// when the caller can resolve the column's type. + pub fn try_new_with_column_type( + predicate: &Expr, + column_type: &impl Fn(&str) -> Option, + ) -> Option { + let columns = predicate.column_refs(); + if columns.len() != 1 { + return None; + } + let column_name = columns.into_iter().next()?.name.clone(); + let column_type = column_type(&column_name); + Some(Self { + column_name, + predicate: SimplePredicate::try_new(predicate, column_type.as_ref())?, + }) + } + + /// Returns the referenced column name. pub fn column_name(&self) -> &str { &self.column_name } pub fn is_eq(&self) -> bool { - matches!(self.op, Operator::Eq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::Eq, + .. + } + ) } - pub fn is_not_eq(&self) -> bool { - matches!(self.op, Operator::NotEq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::NotEq, + .. + } + ) } - pub fn is_lt(&self) -> bool { - matches!(self.op, Operator::Lt) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::Lt, + .. + } + ) } - pub fn is_lt_eq(&self) -> bool { - matches!(self.op, Operator::LtEq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::LtEq, + .. + } + ) } - pub fn is_gt(&self) -> bool { - matches!(self.op, Operator::Gt) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::Gt, + .. + } + ) } - pub fn is_gt_eq(&self) -> bool { - matches!(self.op, Operator::GtEq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::GtEq, + .. + } + ) } - /// Returns true if this filter represents an `OR` chain of equality comparisons, e.g. - /// `col = lit1 OR col = lit2 ...`. + /// Returns true for an OR chain of equality comparisons. pub fn is_or_eq_chain(&self) -> bool { - matches!(self.op, Operator::Or) + matches!(self.predicate, SimplePredicate::OrEqChain { .. }) } - /// Returns the literal as a [`Value`]. It returns `None` if the literal can't be converted. + /// Returns the literal of a bare comparison, without discarding compound conditions. pub fn literal_value(&self) -> Option { - let array = self.literal.get().0; - let scalar = ScalarValue::try_from_array(array, 0).ok()?; - Value::try_from(scalar).ok() + let SimplePredicate::Comparison { literal, .. } = &self.predicate else { + return None; + }; + Value::try_from(ScalarValue::try_from_array(literal.get().0, 0).ok()?).ok() } - /// Returns the literal list as a list of [`Value`]s. It returns `None` if any literal can't be - /// converted. + /// Returns equality-OR literals, or None for other predicates or unsupported values. pub fn literal_list_values(&self) -> Option> { - self.literal_list + let SimplePredicate::OrEqChain { literals } = &self.predicate else { + return None; + }; + literals .iter() - .map(|scalar| { - let array = scalar.get().0; - let scalar = ScalarValue::try_from_array(array, 0).ok()?; - Value::try_from(scalar).ok() + .map(|literal| { + Value::try_from(ScalarValue::try_from_array(literal.get().0, 0).ok()?).ok() }) .collect() } @@ -229,8 +188,7 @@ impl SimpleFilterEvaluator { let input = input .to_scalar() .with_context(|_| ToArrowScalarSnafu { v: input.clone() })?; - let result = self.evaluate_datum(&input, 1)?; - Ok(result.value(0)) + Ok(self.evaluate_datum(&input, 1)?.value(0)) } pub fn evaluate_array(&self, input: &ArrayRef) -> Result { @@ -238,200 +196,27 @@ impl SimpleFilterEvaluator { } pub fn evaluate_vector(&self, input: &VectorRef) -> Result { - self.evaluate_datum(&input.to_arrow_array(), input.len()) + self.evaluate_array(&input.to_arrow_array()) } - fn evaluate_datum(&self, input: &impl Datum, input_len: usize) -> Result { - let result = match self.op { - Operator::Eq => cmp::eq(input, &self.literal), - Operator::NotEq => cmp::neq(input, &self.literal), - Operator::Lt => cmp::lt(input, &self.literal), - Operator::LtEq => cmp::lt_eq(input, &self.literal), - Operator::Gt => cmp::gt(input, &self.literal), - Operator::GtEq => cmp::gt_eq(input, &self.literal), - Operator::RegexMatch => self.regex_match(input), - Operator::RegexIMatch => self.regex_match(input), - Operator::RegexNotMatch => self.regex_match(input), - Operator::RegexNotIMatch => self.regex_match(input), - Operator::Or => { - // OR operator stands for OR-chained EQs (or INLIST in other words) - let mut result: BooleanArray = vec![false; input_len].into(); - for literal in &self.literal_list { - let rhs = cmp::eq(input, literal).context(ArrowComputeSnafu)?; - result = or_kleene(&result, &rhs).context(ArrowComputeSnafu)?; - } - Ok(result) - } - _ => { - return UnsupportedOperationSnafu { - reason: format!("{:?}", self.op), - } - .fail(); - } - }; - result - .context(ArrowComputeSnafu) - .map(|array| boolean_array_to_scan_mask(&array).values().clone()) + fn evaluate_datum(&self, input: &impl Datum, len: usize) -> Result { + let result = self.predicate.evaluate(input, len)?; + Ok(boolean_array_to_scan_mask(&result).values().clone()) } - /// Builds a regex pattern from a scalar value and operator. - /// Returns the `(regex, negative)` and if successful. - /// - /// Returns `Err` if - /// - the value is not a string - /// - the regex pattern is invalid - /// - /// The regex is `None` if - /// - the operator is not a regex operator - /// - the pattern is empty - fn maybe_build_regex( - operator: Operator, - value: &ScalarValue, - ) -> Result<(Option, bool), ArrowError> { - let (ignore_case, negative) = match operator { - Operator::RegexMatch => (false, false), - Operator::RegexIMatch => (true, false), - Operator::RegexNotMatch => (false, true), - Operator::RegexNotIMatch => (true, true), - _ => return Ok((None, false)), - }; - let flag = if ignore_case { Some("i") } else { None }; - let regex = value - .try_as_str() - .ok_or_else(|| ArrowError::CastError(format!("Cannot cast {:?} to str", value)))? - .ok_or_else(|| ArrowError::CastError("Regex should not be null".to_string()))?; - let pattern = match flag { - Some(flag) => format!("(?{flag}){regex}"), - None => regex.to_string(), - }; - if pattern.is_empty() { - Ok((None, negative)) - } else { - Regex::new(pattern.as_str()) - .map_err(|e| { - ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}")) - }) - .map(|regex| (Some(regex), negative)) - } - } - - fn regex_match(&self, input: &impl Datum) -> std::result::Result { - let array = input.get().0; - - // Try to cast to StringArray first - if let Ok(string_array) = as_string_array(array) { - let mut result = regexp_is_match_scalar(string_array, self.regex.as_ref())?; - if self.regex_negative { - result = datatypes::compute::not(&result)?; - } - return Ok(result); - } - - // Try to cast to StringDictionaryArray - if let Some(dict_array) = array.as_any().downcast_ref::>() { - let mut result = regexp_is_match_dictionary(dict_array, self.regex.as_ref())?; - if self.regex_negative { - result = datatypes::compute::not(&result)?; - } - return Ok(result); - } - - Err(ArrowError::CastError(format!( - "Cannot cast {:?} to StringArray or StringDictionaryArray", - array.data_type() - ))) - } - - /// Casts the filter's timestamp literal into the unit of the `target` - /// timestamp type, so it can be evaluated against a column stored in that - /// unit (e.g. an old-unit SST after the time index unit was widened). - /// - /// When the literal is not representable in `target`'s unit (e.g. - /// `ts = 7_000_500us` against a millisecond column), the outcome keeps the - /// row set instead of rounding the literal: `=` prunes / `!=` matches, and - /// inequalities strengthen the operator (e.g. `>= 2_500_500us` becomes - /// `> 2500ms`) so the excluded boundary row stays excluded. - /// - /// Returns `None` when the filter doesn't compare against a tz-naive - /// timestamp literal, or `target` is not a timestamp type. + /// Adapts timestamp predicates to an older SST's timestamp unit without rounding + /// away boundary rows. Nonrepresentable equality literals cannot match a row. pub fn cast_timestamp_unit(&self, target: &ConcreteDataType) -> Option { - let target_unit = match target.as_arrow_type() { - DataType::Timestamp(unit, _) => unit, - _ => return None, + let DataType::Timestamp(target_unit, _) = target.as_arrow_type() else { + return None; }; - - // An `OR` chain of equalities (an IN list): convert each literal - // with `=` semantics; literals not representable in the target unit - // cannot match any row and simply drop out of the chain. - if self.op == Operator::Or { - let mut literal_list = Vec::with_capacity(self.literal_list.len()); - for literal in &self.literal_list { - let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?; - let Some((value, unit)) = timestamp_scalar_parts(&scalar) else { - continue; - }; - let Some(value) = value else { - continue; - }; - let cast = div_mod_units(value, unit.into(), target_unit.into())?; - if cast.remainder == 0 { - literal_list.push(timestamp_scalar(cast.quotient, target_unit)?); - } - } - if literal_list.is_empty() { - return Some(TimestampUnitCast::Pruned); - } - let literal = literal_list[0].clone(); - return Some(TimestampUnitCast::Filter(Self { + Some(match self.predicate.cast_timestamp_unit(target_unit)? { + SimplePredicate::Constant(Some(true)) => TimestampUnitCast::Matched, + SimplePredicate::Constant(_) => TimestampUnitCast::Pruned, + predicate => TimestampUnitCast::Filter(Self { column_name: self.column_name.clone(), - literal, - op: Operator::Or, - literal_list, - regex: None, - regex_negative: false, - })); - } - - let scalar = ScalarValue::try_from_array(self.literal.get().0, 0).ok()?; - let (value, unit) = timestamp_scalar_parts(&scalar)?; - let Some(value) = value else { - // A NULL literal never compares true (arrow comparison - // semantics: null results are filtered out). - return Some(TimestampUnitCast::Pruned); - }; - if unit == target_unit { - return Some(TimestampUnitCast::Filter(self.clone())); - } - let cast = div_mod_units(value, unit.into(), target_unit.into())?; - let divisible = cast.remainder == 0; - let literal = timestamp_scalar(cast.quotient, target_unit)?; - let filter = |op: Operator| Self { - column_name: self.column_name.clone(), - literal: literal.clone(), - op, - literal_list: vec![], - regex: None, - regex_negative: false, - }; - - Some(match self.op { - Operator::Eq if divisible => TimestampUnitCast::Filter(filter(Operator::Eq)), - Operator::Eq => TimestampUnitCast::Pruned, - Operator::NotEq if divisible => TimestampUnitCast::Filter(filter(Operator::NotEq)), - Operator::NotEq => TimestampUnitCast::Matched, - // v > L holds exactly for v > quotient, whether or not L is - // representable in the target unit. - Operator::Gt => TimestampUnitCast::Filter(filter(Operator::Gt)), - Operator::GtEq if divisible => TimestampUnitCast::Filter(filter(Operator::GtEq)), - // L strictly between two target values: v >= L is v > quotient. - Operator::GtEq => TimestampUnitCast::Filter(filter(Operator::Gt)), - Operator::Lt if divisible => TimestampUnitCast::Filter(filter(Operator::Lt)), - // L strictly between two target values: v < L is v <= quotient. - Operator::Lt => TimestampUnitCast::Filter(filter(Operator::LtEq)), - // v <= L holds exactly for v <= quotient. - Operator::LtEq => TimestampUnitCast::Filter(filter(Operator::LtEq)), - // Regex predicates don't apply to timestamps. - _ => return None, + predicate, + }), }) } } @@ -449,35 +234,6 @@ pub enum TimestampUnitCast { Matched, } -/// Extracts the value and unit from a tz-naive timestamp scalar. -fn timestamp_scalar_parts(scalar: &ScalarValue) -> Option<(Option, TimeUnit)> { - let (value, unit, timezone) = match scalar { - ScalarValue::TimestampSecond(v, tz) => (*v, TimeUnit::Second, tz), - ScalarValue::TimestampMillisecond(v, tz) => (*v, TimeUnit::Millisecond, tz), - ScalarValue::TimestampMicrosecond(v, tz) => (*v, TimeUnit::Microsecond, tz), - ScalarValue::TimestampNanosecond(v, tz) => (*v, TimeUnit::Nanosecond, tz), - _ => return None, - }; - // A timezone-aware literal doesn't compare against a tz-naive column; - // leave it to the caller instead of guessing the intended semantics. - (timezone.is_none()).then_some((value, unit)) -} - -/// Builds a tz-naive timestamp literal for `value` in `unit`. -pub fn timestamp_scalar_value(value: i64, unit: TimeUnit) -> ScalarValue { - match unit { - TimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None), - TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None), - TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None), - TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None), - } -} - -/// Builds a one-element scalar array holding `value` in `unit`. -fn timestamp_scalar(value: i64, unit: TimeUnit) -> Option> { - timestamp_scalar_value(value, unit).to_scalar().ok() -} - /// Evaluate the predicate on the input [RecordBatch], and return a new [RecordBatch]. /// Copy from datafusion::physical_plan::src::filter.rs pub fn batch_filter( @@ -520,726 +276,3 @@ fn boolean_array_to_scan_mask(array: &BooleanArray) -> BooleanArray { } BooleanArray::new(values.into(), None) } - -/// The same as arrow [regexp_is_match_scalar()](datatypes::compute::kernels::regexp::regexp_is_match_scalar()) -/// with pre-compiled regex. -/// See for the implementation details. -pub fn regexp_is_match_scalar<'a, S>( - array: &'a S, - regex: Option<&Regex>, -) -> Result -where - &'a S: StringArrayType<'a>, -{ - let null_bit_buffer = array.nulls().map(|x| x.inner().sliced()); - let mut result = BooleanBufferBuilder::new(array.len()); - - if let Some(re) = regex { - for i in 0..array.len() { - let value = array.value(i); - result.append(re.is_match(value)); - } - } else { - result.append_n(array.len(), true); - } - - let buffer = result.into(); - let data = unsafe { - ArrayData::new_unchecked( - DataType::Boolean, - array.len(), - None, - null_bit_buffer, - 0, - vec![buffer], - vec![], - ) - }; - - Ok(BooleanArray::from(data)) -} - -/// Similar to [regexp_is_match_scalar] but for StringDictionaryArray. -/// Iterates through dictionary keys to get string values and applies regex matching. -pub fn regexp_is_match_dictionary( - dict_array: &DictionaryArray, - regex: Option<&Regex>, -) -> Result { - // Get the string values from the dictionary - let string_values = dict_array - .values() - .as_any() - .downcast_ref::() - .ok_or_else(|| { - ArrowError::CastError("Dictionary values must be StringArray".to_string()) - })?; - - // Dictionary logical nulls include both null keys and keys whose dictionary value is null. - let logical_nulls = dict_array.logical_nulls(); - let null_bit_buffer = logical_nulls.as_ref().map(|x| x.inner().sliced()); - let mut result = BooleanBufferBuilder::new(dict_array.len()); - - if let Some(re) = regex { - let keys = dict_array.keys().values(); - for i in 0..dict_array.len() { - if logical_nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) { - result.append(false); - } else { - let key = keys[i] as usize; - let string_value = string_values.value(key); - result.append(re.is_match(string_value)); - } - } - } else { - result.append_n(dict_array.len(), true); - } - - let buffer = result.into(); - let data = unsafe { - ArrayData::new_unchecked( - DataType::Boolean, - dict_array.len(), - None, - null_bit_buffer, - 0, - vec![buffer], - vec![], - ) - }; - - Ok(BooleanArray::from(data)) -} - -#[cfg(test)] -mod test { - - use std::sync::Arc; - - use datafusion::execution::context::ExecutionProps; - use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; - use datafusion::logical_expr::{BinaryExpr, col, lit}; - use datafusion::physical_expr::create_physical_expr; - use datafusion_common::{Column, DFSchema}; - use datatypes::arrow::array::{TimestampMillisecondArray, TimestampNanosecondArray}; - use datatypes::arrow::datatypes::{DataType, Field, Schema}; - - use super::*; - - #[test] - fn unsupported_filter_op() { - // `+` is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Plus, - right: Box::new(1.lit()), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - - // two literal is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(1.lit()), - op: Operator::Eq, - right: Box::new(1.lit()), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - - // two column is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(Expr::Column(Column::from_name("bar"))), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - - // compound expr is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(1.lit()), - })), - op: Operator::Eq, - right: Box::new(1.lit()), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - } - - #[test] - fn supported_filter_op() { - // equal - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(1.lit()), - }); - let _ = SimpleFilterEvaluator::try_new(&expr).unwrap(); - - // swap operands - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(1.lit()), - op: Operator::Lt, - right: Box::new(Expr::Column(Column::from_name("foo"))), - }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); - assert_eq!(evaluator.op, Operator::Gt); - assert_eq!(evaluator.column_name, "foo".to_string()); - } - - #[test] - fn run_on_array() { - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(1i64.lit()), - }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); - - let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _; - let result = evaluator.evaluate_array(&input_1).unwrap(); - assert_eq!(result, BooleanBuffer::from(vec![true, false, false])); - - let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _; - let result = evaluator.evaluate_array(&input_2).unwrap(); - assert_eq!(result, BooleanBuffer::from(vec![true, true, true])); - - let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _; - let result = evaluator.evaluate_array(&input_3).unwrap(); - assert_eq!(result, BooleanBuffer::from(vec![])); - } - - #[test] - fn run_on_scalar() { - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Lt, - right: Box::new(1i64.lit()), - }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); - - let input_1 = ScalarValue::Int64(Some(1)); - let result = evaluator.evaluate_scalar(&input_1).unwrap(); - assert!(!result); - - let input_2 = ScalarValue::Int64(Some(0)); - let result = evaluator.evaluate_scalar(&input_2).unwrap(); - assert!(result); - - let input_3 = ScalarValue::Int64(None); - let result = evaluator.evaluate_scalar(&input_3).unwrap(); - assert!(!result); - } - - #[test] - fn batch_filter_test() { - let expr = col("ts").gt(lit(123456u64)); - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("ts", DataType::UInt64, false), - ]); - let df_schema = DFSchema::try_from(schema.clone()).unwrap(); - let props = ExecutionProps::new(); - let physical_expr = create_physical_expr( - &expr, - &df_schema, - &props, - &PhysicalPlanningContext::default(), - ) - .unwrap(); - let batch = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(datatypes::arrow::array::Int32Array::from(vec![4, 5, 6])), - Arc::new(datatypes::arrow::array::UInt64Array::from(vec![ - 123456, 123457, 123458, - ])), - ], - ) - .unwrap(); - let new_batch = batch_filter(&batch, &physical_expr).unwrap(); - assert_eq!(new_batch.num_rows(), 2); - let first_column_values = new_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let expected = datatypes::arrow::array::Int32Array::from(vec![5, 6]); - assert_eq!(first_column_values, &expected); - } - - #[test] - fn test_complex_filter_expression() { - // Create an expression tree for: col = 'B' OR col = 'C' OR col = 'D' - let col_eq_b = col("col").eq(lit("B")); - let col_eq_c = col("col").eq(lit("C")); - let col_eq_d = col("col").eq(lit("D")); - - // Build the OR chain - let col_or_expr = col_eq_b.or(col_eq_c).or(col_eq_d); - - // Check that SimpleFilterEvaluator can handle OR chain - let or_evaluator = SimpleFilterEvaluator::try_new(&col_or_expr).unwrap(); - assert_eq!(or_evaluator.column_name, "col"); - assert_eq!(or_evaluator.op, Operator::Or); - assert_eq!(or_evaluator.literal_list.len(), 3); - assert_eq!( - format!("{:?}", or_evaluator.literal_list), - "[Scalar(StringArray\n[\n \"B\",\n]), Scalar(StringArray\n[\n \"C\",\n]), Scalar(StringArray\n[\n \"D\",\n])]" - ); - - // Create a schema and batch for testing - let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]); - let df_schema = DFSchema::try_from(schema.clone()).unwrap(); - let props = ExecutionProps::new(); - let physical_expr = create_physical_expr( - &col_or_expr, - &df_schema, - &props, - &PhysicalPlanningContext::default(), - ) - .unwrap(); - - // Create test data - let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![ - "B", "C", "E", "B", "C", "D", "F", - ])); - let batch = RecordBatch::try_new(Arc::new(schema), vec![col_data]).unwrap(); - let expected = datatypes::arrow::array::StringArray::from(vec!["B", "C", "B", "C", "D"]); - - // Filter the batch - let filtered_batch = batch_filter(&batch, &physical_expr).unwrap(); - - // Expected: rows with col in ("B", "C", "D") - // That would be rows 0, 1, 3, 4, 5 - assert_eq!(filtered_batch.num_rows(), 5); - - let col_filtered = filtered_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(col_filtered, &expected); - } - - #[test] - fn test_maybe_build_regex() { - // Test case for RegexMatch (case sensitive, non-negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(!negative); - assert!(regex.unwrap().is_match("axxb")); - - // Test case for RegexIMatch (case insensitive, non-negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexIMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(!negative); - assert!(regex.unwrap().is_match("AxxB")); - - // Test case for RegexNotMatch (case sensitive, negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexNotMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(negative); - - // Test case for RegexNotIMatch (case insensitive, negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexNotIMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(negative); - - // Test with empty regex pattern - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(Some("".to_string())), - ) - .unwrap(); - assert!(regex.is_none()); - assert!(!negative); - - // Test with non-regex operator - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::Eq, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_none()); - assert!(!negative); - - // Test with invalid regex pattern - let result = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(Some("a(b".to_string())), - ); - assert!(result.is_err()); - - // Test with non-string value - let result = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Int64(Some(123)), - ); - assert!(result.is_err()); - - // Test with null value - let result = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(None), - ); - assert!(result.is_err()); - } - - #[test] - fn test_regex_match_dictionary_array() { - use datatypes::arrow::array::StringDictionaryBuilder; - - // Create a StringDictionaryArray - let mut builder = StringDictionaryBuilder::::new(); - builder.append("apple").unwrap(); - builder.append("banana").unwrap(); - builder.append("apple").unwrap(); - builder.append("cherry").unwrap(); - let dict_array = builder.finish(); - - // Test regex that matches "apple" - let regex = regex::Regex::new(r"app.*").unwrap(); - let result = regexp_is_match_dictionary(&dict_array, Some(®ex)).unwrap(); - - // Should match indices 0 and 2 (both "apple") - assert_eq!(result.len(), 4); - assert!(result.value(0)); // "apple" - assert!(!result.value(1)); // "banana" - assert!(result.value(2)); // "apple" - assert!(!result.value(3)); // "cherry" - - // Test regex that matches "banana" - let regex2 = regex::Regex::new(r"ban.*").unwrap(); - let result2 = regexp_is_match_dictionary(&dict_array, Some(®ex2)).unwrap(); - - assert!(!result2.value(0)); // "apple" - assert!(result2.value(1)); // "banana" - assert!(!result2.value(2)); // "apple" - assert!(!result2.value(3)); // "cherry" - - // Test with no regex (should match all) - let result3 = regexp_is_match_dictionary(&dict_array, None).unwrap(); - assert!(result3.value(0)); - assert!(result3.value(1)); - assert!(result3.value(2)); - assert!(result3.value(3)); - } - - #[test] - fn test_regex_scan_masks_preserve_sql_null_semantics() { - let plain = Arc::new(datatypes::arrow::array::StringArray::from(vec![ - Some("api"), - Some("API"), - Some("db"), - None, - ])) as ArrayRef; - assert_regex_scan_masks( - &plain, - [ - vec![true, false, false, false], - vec![true, true, false, false], - vec![false, true, true, false], - vec![false, false, true, false], - ], - ); - - let dictionary = DictionaryArray::new( - datatypes::arrow::array::UInt32Array::from(vec![ - Some(0), - Some(1), - Some(2), - None, - Some(3), - ]), - Arc::new(datatypes::arrow::array::StringArray::from(vec![ - Some("api"), - Some("API"), - Some("db"), - None, - ])), - ); - let raw = - regexp_is_match_dictionary(&dictionary, Some(&Regex::new("^api$").unwrap())).unwrap(); - assert!(raw.is_null(3)); // null dictionary key - assert!(raw.is_null(4)); // non-null key referencing a null dictionary value - - let dictionary = Arc::new(dictionary) as ArrayRef; - assert_regex_scan_masks( - &dictionary, - [ - vec![true, false, false, false, false], - vec![true, true, false, false, false], - vec![false, true, true, false, false], - vec![false, false, true, false, false], - ], - ); - - let negative = regex_evaluator(Operator::RegexNotMatch); - assert!(!negative.evaluate_scalar(&ScalarValue::Utf8(None)).unwrap()); - } - - #[test] - fn test_nullable_boolean_predicate_becomes_scan_mask() { - let predicate = BooleanArray::from(vec![Some(true), None, Some(false)]); - assert_eq!( - BooleanArray::from(vec![true, false, false]), - boolean_array_to_scan_mask(&predicate) - ); - } - - fn assert_regex_scan_masks(input: &ArrayRef, expected: [Vec; 4]) { - for (op, expected) in [ - Operator::RegexMatch, - Operator::RegexIMatch, - Operator::RegexNotMatch, - Operator::RegexNotIMatch, - ] - .into_iter() - .zip(expected) - { - assert_eq!( - BooleanBuffer::from(expected), - regex_evaluator(op).evaluate_array(input).unwrap(), - "{op:?}" - ); - } - } - - fn regex_evaluator(op: Operator) -> SimpleFilterEvaluator { - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("host"))), - op, - right: Box::new("^api$".lit()), - }); - SimpleFilterEvaluator::try_new(&expr).unwrap() - } - - fn ts_us(v: i64) -> ScalarValue { - ScalarValue::TimestampMicrosecond(Some(v), None) - } - - fn cast_to_ms(expr: &Expr) -> Option { - SimpleFilterEvaluator::try_new(expr) - .unwrap() - .cast_timestamp_unit(&ConcreteDataType::timestamp_millisecond_datatype()) - } - - /// Evaluates a cast `Filter` outcome against millisecond values and - /// returns the mask. - fn eval_ms_mask(cast: Option, values: &[i64]) -> Vec { - let filter = match cast.expect("cast must apply") { - TimestampUnitCast::Filter(filter) => filter, - other => panic!("expected Filter outcome, got {other:?}"), - }; - let array = Arc::new(TimestampMillisecondArray::from(values.to_vec())) as ArrayRef; - filter.evaluate_array(&array).unwrap().iter().collect() - } - - #[test] - fn cast_timestamp_unit_converts_representable_literal() { - // ts = 7_000_000us evaluated against a ms column becomes ts = 7000ms. - assert_eq!( - vec![false, true, false], - eval_ms_mask( - cast_to_ms(&col("ts").eq(lit(ts_us(7_000_000)))), - &[6_999, 7_000, 7_001] - ) - ); - // != complements =. - assert_eq!( - vec![true, false, true], - eval_ms_mask( - cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_000)))), - &[6_999, 7_000, 7_001] - ) - ); - } - - #[test] - fn cast_timestamp_unit_prunes_non_representable_equality() { - // 7_000_500us is not a whole millisecond: no ms row can equal it, - // and every ms row satisfies `!=`. - assert!(matches!( - cast_to_ms(&col("ts").eq(lit(ts_us(7_000_500)))), - Some(TimestampUnitCast::Pruned) - )); - assert!(matches!( - cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_500)))), - Some(TimestampUnitCast::Matched) - )); - } - - #[test] - fn cast_timestamp_unit_strengthens_inequalities() { - // 2_500_500us is strictly between 2500ms and 2501ms: the cast must - // not round the literal to a boundary the original predicate - // excludes (>= 2500ms would wrongly match 2500ms). - let values = vec![2_499, 2_500, 2_501]; - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_500)))), &values) - ); - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_500)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_500)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_500)))), &values) - ); - - // Representable boundary keeps the original operator. - let values = vec![2_499, 2_500, 2_501]; - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_000)))), &values) - ); - assert_eq!( - vec![false, true, true], - eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_000)))), &values) - ); - assert_eq!( - vec![true, false, false], - eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_000)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_000)))), &values) - ); - } - - #[test] - fn cast_timestamp_unit_handles_negative_instants() { - // -2_500_500us floors to -2501ms with remainder 500us: only rows - // with instant >= -2.5005ms may match >= / >. - let values = vec![-2_502, -2_501, -2_500]; - assert_eq!( - vec![false, false, true], - eval_ms_mask( - cast_to_ms(&col("ts").gt_eq(lit(ts_us(-2_500_500)))), - &values - ) - ); - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(-2_500_500)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask( - cast_to_ms(&col("ts").lt_eq(lit(ts_us(-2_500_500)))), - &values - ) - ); - // Exactly representable negative literal: -2_500_000us == -2500ms. - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").eq(lit(ts_us(-2_500_000)))), &values) - ); - } - - #[test] - fn cast_timestamp_unit_or_chain_drops_unrepresentable_literals() { - // Only 7_000_000us is a whole millisecond; 7_000_500us cannot match. - let expr = col("ts") - .eq(lit(ts_us(7_000_500))) - .or(col("ts").eq(lit(ts_us(7_000_000)))); - assert_eq!( - vec![false, true], - eval_ms_mask(cast_to_ms(&expr), &[6_999, 7_000]) - ); - - // A chain where no literal is representable matches nothing. - let expr = col("ts") - .eq(lit(ts_us(7_000_500))) - .or(col("ts").eq(lit(ts_us(6_000_500)))); - assert!(matches!(cast_to_ms(&expr), Some(TimestampUnitCast::Pruned))); - } - - #[test] - fn cast_timestamp_unit_null_literal_prunes() { - // A NULL literal never compares true. - assert!(matches!( - cast_to_ms(&col("ts").eq(lit(ScalarValue::TimestampMicrosecond(None, None)))), - Some(TimestampUnitCast::Pruned) - )); - } - - #[test] - fn cast_timestamp_unit_same_unit_returns_filter_directly() { - // A literal already in the target unit is returned unchanged. - let expr = col("ts").eq(lit(ts_us(7_000_000))); - let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); - let cast = filter - .cast_timestamp_unit(&ConcreteDataType::timestamp_microsecond_datatype()) - .unwrap(); - let TimestampUnitCast::Filter(f) = cast else { - panic!("expected Filter, got {cast:?}") - }; - assert_eq!(filter.op, f.op); - assert_eq!(filter.column_name(), f.column_name()); - assert_eq!( - ts_us(7_000_000), - ScalarValue::try_from_array(f.literal.get().0, 0).unwrap() - ); - } - - #[test] - fn cast_timestamp_unit_rejects_non_timestamps() { - // Non-timestamp target type. - assert!( - SimpleFilterEvaluator::try_new(&col("ts").gt(lit(ts_us(1)))) - .unwrap() - .cast_timestamp_unit(&ConcreteDataType::int64_datatype()) - .is_none() - ); - // Non-timestamp literal against a timestamp target. - assert!(cast_to_ms(&col("ts").gt(lit(42_i64))).is_none()); - } - - #[test] - fn cast_timestamp_unit_to_finer_unit() { - // 5 seconds against a nanosecond column becomes 5_000_000_000ns. - let filter = SimpleFilterEvaluator::try_new( - &col("ts").eq(lit(ScalarValue::TimestampSecond(Some(5), None))), - ) - .unwrap() - .cast_timestamp_unit(&ConcreteDataType::timestamp_nanosecond_datatype()) - .and_then(|cast| match cast { - TimestampUnitCast::Filter(filter) => Some(filter), - _ => None, - }) - .unwrap(); - let array = Arc::new(TimestampNanosecondArray::from(vec![ - 4_999_999_999, - 5_000_000_000, - ])) as ArrayRef; - assert_eq!( - vec![false, true], - filter - .evaluate_array(&array) - .unwrap() - .iter() - .collect::>() - ); - } -} diff --git a/src/common/recordbatch/src/filter/predicate.rs b/src/common/recordbatch/src/filter/predicate.rs new file mode 100644 index 00000000000..e283633dc02 --- /dev/null +++ b/src/common/recordbatch/src/filter/predicate.rs @@ -0,0 +1,944 @@ +// 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. + +//! Predicate recognition, execution, and precompiled regular expressions. + +use datafusion::logical_expr::{Expr, Operator}; +use datafusion_common::ScalarValue; +use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar}; +use datafusion_common::arrow::buffer::BooleanBuffer; +use datafusion_common::arrow::compute::kernels::cmp; +use datafusion_common::cast::as_string_array; +use datatypes::arrow::array::{ + Array, ArrayAccessor, ArrayData, BooleanArray, BooleanBufferBuilder, DictionaryArray, + StringArrayType, +}; +use datatypes::arrow::datatypes::{DataType, UInt32Type}; +use datatypes::arrow::error::ArrowError; +use datatypes::compute::or_kleene; +use regex::Regex; +use snafu::ResultExt; + +use crate::error::{ArrowComputeSnafu, Result, UnsupportedOperationSnafu}; + +#[derive(Debug, Clone)] +pub(crate) enum SimplePredicate { + Comparison { + literal: Scalar, + op: Operator, + regex: Option, + regex_negative: bool, + }, + OrEqChain { + literals: Vec>, + }, + IsNull { + negated: bool, + }, + Constant(Option), + And(Box, Box), + Or(Box, Box), +} + +impl SimplePredicate { + pub(crate) fn try_new(expr: &Expr, column_type: Option<&DataType>) -> Option { + match expr { + Expr::BinaryExpr(binary) if matches!(binary.op, Operator::And | Operator::Or) => { + let left = Self::try_new(&binary.left, column_type)?; + let right = Self::try_new(&binary.right, column_type)?; + if binary.op == Operator::And { + return Some(Self::And(Box::new(left), Box::new(right))); + } + // Preserve the encoded-key IN-list fast path for pure equality disjunctions. + if let (Some(mut lhs), Some(rhs)) = + (left.equality_literals(), right.equality_literals()) + { + lhs.extend(rhs); + return Some(Self::OrEqChain { literals: lhs }); + } + Some(Self::Or(Box::new(left), Box::new(right))) + } + Expr::BinaryExpr(binary) => { + let mut op = binary.op; + if !matches!( + op, + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + | Operator::RegexMatch + | Operator::RegexIMatch + | Operator::RegexNotMatch + | Operator::RegexNotIMatch + ) { + return None; + } + let literal = if let Expr::Literal(literal, _) = &*binary.right { + Self::check_column(&binary.left, column_type)?; + literal + } else if let Expr::Literal(literal, _) = &*binary.left { + Self::check_column(&binary.right, column_type)?; + op = op.swap()?; + literal + } else { + return None; + }; + let (regex, regex_negative) = Self::maybe_build_regex(op, literal).ok()?; + Some(Self::Comparison { + literal: literal.to_scalar().ok()?, + op, + regex, + regex_negative, + }) + } + Expr::IsNull(column) | Expr::IsNotNull(column) => { + Self::check_column(column, column_type)?; + Some(Self::IsNull { + negated: matches!(expr, Expr::IsNotNull(_)), + }) + } + Expr::Literal(ScalarValue::Boolean(value), _) => Some(Self::Constant(*value)), + _ => None, + } + } + + fn check_column(expr: &Expr, column_type: Option<&DataType>) -> Option<()> { + match expr { + Expr::Column(_) => Some(()), + Expr::Cast(cast) if cast.field.data_type() == &DataType::Utf8 => { + let Expr::Column(_) = &*cast.expr else { + return None; + }; + match column_type? { + DataType::Utf8 => Some(()), + DataType::Dictionary(key, value) + if **key == DataType::UInt32 && **value == DataType::Utf8 => + { + Some(()) + } + _ => None, + } + } + _ => None, + } + } + + fn equality_literals(&self) -> Option>> { + match self { + Self::Comparison { + literal, + op: Operator::Eq, + .. + } => Some(vec![literal.clone()]), + Self::OrEqChain { literals } => Some(literals.clone()), + _ => None, + } + } + + pub(crate) fn evaluate(&self, input: &impl Datum, len: usize) -> Result { + match self { + Self::Comparison { + literal, + op, + regex, + regex_negative, + } => match op { + Operator::Eq => cmp::eq(input, literal), + Operator::NotEq => cmp::neq(input, literal), + Operator::Lt => cmp::lt(input, literal), + Operator::LtEq => cmp::lt_eq(input, literal), + Operator::Gt => cmp::gt(input, literal), + Operator::GtEq => cmp::gt_eq(input, literal), + Operator::RegexMatch + | Operator::RegexIMatch + | Operator::RegexNotMatch + | Operator::RegexNotIMatch => { + Self::regex_match(input, regex.as_ref(), *regex_negative) + } + _ => { + return UnsupportedOperationSnafu { + reason: format!("{op:?}"), + } + .fail(); + } + } + .context(ArrowComputeSnafu), + Self::OrEqChain { literals } => { + let mut result = BooleanArray::from(vec![false; len]); + for literal in literals { + let rhs = cmp::eq(input, literal).context(ArrowComputeSnafu)?; + result = or_kleene(&result, &rhs).context(ArrowComputeSnafu)?; + } + Ok(result) + } + Self::IsNull { negated } => { + let values = match input.get().0.logical_nulls() { + Some(nulls) if *negated => nulls.inner().clone(), + Some(nulls) => !nulls.inner(), + None if *negated => BooleanBuffer::new_set(len), + None => BooleanBuffer::new_unset(len), + }; + Ok(BooleanArray::new(values, None)) + } + Self::Constant(value) => Ok(BooleanArray::from(vec![*value; len])), + Self::And(left, right) => datatypes::compute::and_kleene( + &left.evaluate(input, len)?, + &right.evaluate(input, len)?, + ) + .context(ArrowComputeSnafu), + Self::Or(left, right) => { + or_kleene(&left.evaluate(input, len)?, &right.evaluate(input, len)?) + .context(ArrowComputeSnafu) + } + } + } + + /// Builds a regex pattern from a scalar value and operator. + /// Returns the `(regex, negative)` and if successful. + /// + /// Returns `Err` if + /// - the value is not a string + /// - the regex pattern is invalid + /// + /// The regex is `None` if + /// - the operator is not a regex operator + /// - the pattern is empty + fn maybe_build_regex( + operator: Operator, + value: &ScalarValue, + ) -> Result<(Option, bool), ArrowError> { + let (ignore_case, negative) = match operator { + Operator::RegexMatch => (false, false), + Operator::RegexIMatch => (true, false), + Operator::RegexNotMatch => (false, true), + Operator::RegexNotIMatch => (true, true), + _ => return Ok((None, false)), + }; + let flag = if ignore_case { Some("i") } else { None }; + let regex = value + .try_as_str() + .ok_or_else(|| ArrowError::CastError(format!("Cannot cast {:?} to str", value)))? + .ok_or_else(|| ArrowError::CastError("Regex should not be null".to_string()))?; + let pattern = match flag { + Some(flag) => format!("(?{flag}){regex}"), + None => regex.to_string(), + }; + if pattern.is_empty() { + Ok((None, negative)) + } else { + Regex::new(pattern.as_str()) + .map_err(|e| { + ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}")) + }) + .map(|regex| (Some(regex), negative)) + } + } + + fn regex_match( + input: &impl Datum, + regex: Option<&Regex>, + negative: bool, + ) -> std::result::Result { + let array = input.get().0; + + // Try to cast to StringArray first + if let Ok(string_array) = as_string_array(array) { + let mut result = regexp_is_match_scalar(string_array, regex)?; + if negative { + result = datatypes::compute::not(&result)?; + } + return Ok(result); + } + + // Try to cast to StringDictionaryArray + if let Some(dict_array) = array.as_any().downcast_ref::>() { + let mut result = regexp_is_match_dictionary(dict_array, regex)?; + if negative { + result = datatypes::compute::not(&result)?; + } + return Ok(result); + } + + Err(ArrowError::CastError(format!( + "Cannot cast {:?} to StringArray or StringDictionaryArray", + array.data_type() + ))) + } +} + +/// The same as arrow [regexp_is_match_scalar()](datatypes::compute::kernels::regexp::regexp_is_match_scalar()) +/// with pre-compiled regex. +/// See for the implementation details. +pub fn regexp_is_match_scalar<'a, S>( + array: &'a S, + regex: Option<&Regex>, +) -> Result +where + &'a S: StringArrayType<'a>, +{ + let null_bit_buffer = array.nulls().map(|x| x.inner().sliced()); + let mut result = BooleanBufferBuilder::new(array.len()); + + if let Some(re) = regex { + for i in 0..array.len() { + let value = array.value(i); + result.append(re.is_match(value)); + } + } else { + result.append_n(array.len(), true); + } + + let buffer = result.into(); + let data = unsafe { + ArrayData::new_unchecked( + DataType::Boolean, + array.len(), + None, + null_bit_buffer, + 0, + vec![buffer], + vec![], + ) + }; + + Ok(BooleanArray::from(data)) +} + +/// Similar to [regexp_is_match_scalar] but for StringDictionaryArray. +/// Iterates through dictionary keys to get string values and applies regex matching. +pub fn regexp_is_match_dictionary( + dict_array: &DictionaryArray, + regex: Option<&Regex>, +) -> Result { + // Get the string values from the dictionary + let string_values = dict_array + .values() + .as_any() + .downcast_ref::() + .ok_or_else(|| { + ArrowError::CastError("Dictionary values must be StringArray".to_string()) + })?; + + // Dictionary logical nulls include both null keys and keys whose dictionary value is null. + let logical_nulls = dict_array.logical_nulls(); + let null_bit_buffer = logical_nulls.as_ref().map(|x| x.inner().sliced()); + let mut result = BooleanBufferBuilder::new(dict_array.len()); + + if let Some(re) = regex { + let keys = dict_array.keys().values(); + for i in 0..dict_array.len() { + if logical_nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) { + result.append(false); + } else { + let key = keys[i] as usize; + let string_value = string_values.value(key); + result.append(re.is_match(string_value)); + } + } + } else { + result.append_n(dict_array.len(), true); + } + + let buffer = result.into(); + let data = unsafe { + ArrayData::new_unchecked( + DataType::Boolean, + dict_array.len(), + None, + null_bit_buffer, + 0, + vec![buffer], + vec![], + ) + }; + + Ok(BooleanArray::from(data)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::execution::context::ExecutionProps; + use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; + use datafusion::logical_expr::{BinaryExpr, Expr, Literal, Operator, col, lit}; + use datafusion::physical_expr::create_physical_expr; + use datafusion_common::arrow::array::ArrayRef; + use datafusion_common::arrow::buffer::BooleanBuffer; + use datafusion_common::{Column, DFSchema, ScalarValue}; + use datatypes::arrow::array::{Array, DictionaryArray, RecordBatch}; + use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; + use datatypes::value::Value; + use regex::Regex; + + use crate::filter::predicate::{SimplePredicate, regexp_is_match_dictionary}; + use crate::filter::{SimpleFilterEvaluator, batch_filter}; + + #[test] + fn unsupported_filter_op() { + // `+` is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Plus, + right: Box::new(1.lit()), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + + // two literal is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(1.lit()), + op: Operator::Eq, + right: Box::new(1.lit()), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + + // two column is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(Expr::Column(Column::from_name("bar"))), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + + // compound expr is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(1.lit()), + })), + op: Operator::Eq, + right: Box::new(1.lit()), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + for negated in [false, true] { + let membership = col("label").in_list(vec![lit(""), lit("lo")], negated); + assert!(SimpleFilterEvaluator::try_new(&membership).is_none()); + assert!( + SimpleFilterEvaluator::try_new(&col("label").is_null().or(membership)).is_none() + ); + } + for expr in [ + col("other").is_null().or(col("label").not_eq(lit("lo"))), + col("a.label") + .is_null() + .or(col("b.label").not_eq(lit("lo"))), + col("label").is_null().or(col("label").eq(col("other"))), + col("label") + .is_null() + .or(Expr::Cast(datafusion::logical_expr::Cast::new( + Box::new(col("label")), + DataType::Int64, + )) + .eq(lit("lo"))), + ] { + assert!(SimpleFilterEvaluator::try_new(&expr).is_none(), "{expr}"); + } + } + + #[test] + fn string_cast_requires_matching_column_type() { + let cast = Expr::Cast(datafusion::logical_expr::Cast::new( + Box::new(col("label")), + DataType::Utf8, + )); + let expr = col("label").is_null().or(cast.not_eq(lit("lo"))); + for data_type in [ + DataType::Int64, + DataType::Timestamp(TimeUnit::Millisecond, None), + DataType::Binary, + ] { + assert!( + SimpleFilterEvaluator::try_new_with_column_type(&expr, &|_| Some( + data_type.clone() + )) + .is_none() + ); + } + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + } + + #[tokio::test] + async fn single_column_predicates_match_datafusion() { + use datafusion::prelude::SessionContext; + use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + use datatypes::arrow::array::{Int32Array, StringArray, UInt32Array}; + + let values = vec![ + None, + Some(""), + Some("lo"), + Some("eth0"), + Some("a\nb"), + Some("服务"), + None, + ]; + let strings: ArrayRef = Arc::new(StringArray::from(values)); + let dictionary: ArrayRef = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![ + None, + Some(0), + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + ]), + Arc::new(StringArray::from(vec![ + Some(""), + Some("lo"), + Some("eth0"), + Some("a\nb"), + Some("服务"), + None, + ])), + )); + let null = || lit(ScalarValue::Utf8(None)); + let mut predicates = vec![ + col("label").is_null(), + col("label").is_not_null(), + col("label") + .is_null() + .or(col("label").eq(lit(""))) + .and(col("label").is_not_null()), + col("label") + .not_eq(lit("lo")) + .and(col("label").not_eq(lit("eth0"))), + col("label") + .eq(lit("lo")) + .or(col("label").eq(lit("eth0"))) + .or(col("label").eq(lit(""))), + col("label").eq(null()).or(col("label").eq(lit("eth0"))), + col("label").gt(lit("a")).and(col("label").lt(lit("z"))), + lit("lo").not_eq(col("label")), + ]; + for (op, literal) in [ + (Operator::Eq, ""), + (Operator::NotEq, "lo"), + (Operator::RegexNotMatch, "^(?:lo|vfat)$"), + (Operator::RegexMatch, "^(?:|eth0)$"), + ] { + let comparison = Expr::BinaryExpr(BinaryExpr::new( + Box::new(col("label")), + op, + Box::new(lit(literal)), + )); + predicates.extend([ + col("label").is_null().or(comparison.clone()), + comparison.or(col("label").is_null()), + ]); + } + for input in [strings, dictionary] { + let ctx = SessionContext::new(); + let schema = Arc::new(Schema::new(vec![ + Field::new("label", input.data_type().clone(), true), + Field::new("id", DataType::Int32, false), + ])); + ctx.register_batch( + "labels", + RecordBatch::try_new( + schema, + vec![ + input.clone(), + Arc::new(Int32Array::from_iter_values(0..input.len() as i32)), + ], + ) + .unwrap(), + ) + .unwrap(); + for expr in &predicates { + let df = ctx + .table("labels") + .await + .unwrap() + .filter(expr.clone()) + .unwrap(); + let plan = df.clone().into_optimized_plan().unwrap(); + let expected = df + .collect() + .await + .unwrap() + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + let evaluate = |expr: &Expr| { + let evaluator = + SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| { + (name == "label").then(|| input.data_type().clone()) + }) + .unwrap_or_else(|| panic!("unsupported: {expr}")); + let mask = evaluator.evaluate_array(&input).unwrap(); + for (i, matched) in mask.iter().enumerate() { + let scalar = ScalarValue::try_from_array(&input, i).unwrap(); + assert_eq!( + evaluator.evaluate_scalar(&scalar).unwrap(), + matched, + "{expr}, row={i}" + ); + } + mask + }; + let selected = |mask: BooleanBuffer| { + mask.iter() + .enumerate() + .filter_map(|(i, matched)| matched.then_some(i as i32)) + .collect::>() + }; + assert_eq!(selected(evaluate(expr)), expected, "{expr}"); + let mut filters = Vec::new(); + plan.apply(|plan| { + if let datafusion::logical_expr::LogicalPlan::Filter(filter) = plan { + filters.push(filter.predicate.clone()); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + if !filters.is_empty() { + let mut mask = BooleanBuffer::new_set(input.len()); + for filter in filters { + mask = &mask & &evaluate(&filter); + } + assert_eq!(selected(mask), expected, "optimized: {plan}"); + } + } + } + } + + #[test] + fn supported_filter_op() { + // equal + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(1.lit()), + }); + let _ = SimpleFilterEvaluator::try_new(&expr).unwrap(); + + // swap operands + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(1.lit()), + op: Operator::Lt, + right: Box::new(Expr::Column(Column::from_name("foo"))), + }); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + assert!(evaluator.is_gt()); + assert_eq!(evaluator.column_name, "foo".to_string()); + } + + #[test] + fn run_on_array() { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(1i64.lit()), + }); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + + let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _; + let result = evaluator.evaluate_array(&input_1).unwrap(); + assert_eq!(result, BooleanBuffer::from(vec![true, false, false])); + + let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _; + let result = evaluator.evaluate_array(&input_2).unwrap(); + assert_eq!(result, BooleanBuffer::from(vec![true, true, true])); + + let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _; + let result = evaluator.evaluate_array(&input_3).unwrap(); + assert_eq!(result, BooleanBuffer::from(vec![])); + } + + #[test] + fn run_on_scalar() { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Lt, + right: Box::new(1i64.lit()), + }); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + + let input_1 = ScalarValue::Int64(Some(1)); + let result = evaluator.evaluate_scalar(&input_1).unwrap(); + assert!(!result); + + let input_2 = ScalarValue::Int64(Some(0)); + let result = evaluator.evaluate_scalar(&input_2).unwrap(); + assert!(result); + + let input_3 = ScalarValue::Int64(None); + let result = evaluator.evaluate_scalar(&input_3).unwrap(); + assert!(!result); + } + + #[test] + fn test_complex_filter_expression() { + // Create an expression tree for: col = 'B' OR col = 'C' OR col = 'D' + let col_eq_b = col("col").eq(lit("B")); + let col_eq_c = col("col").eq(lit("C")); + let col_eq_d = col("col").eq(lit("D")); + + // Build the OR chain + let col_or_expr = col_eq_b.or(col_eq_c).or(col_eq_d); + + // Check that SimpleFilterEvaluator can handle OR chain + let or_evaluator = SimpleFilterEvaluator::try_new(&col_or_expr).unwrap(); + assert_eq!(or_evaluator.column_name, "col"); + assert!(or_evaluator.is_or_eq_chain()); + assert_eq!( + or_evaluator.literal_list_values().unwrap(), + vec![Value::from("B"), Value::from("C"), Value::from("D")] + ); + + // Create a schema and batch for testing + let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]); + let df_schema = DFSchema::try_from(schema.clone()).unwrap(); + let props = ExecutionProps::new(); + let physical_expr = create_physical_expr( + &col_or_expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); + + // Create test data + let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![ + "B", "C", "E", "B", "C", "D", "F", + ])); + let batch = RecordBatch::try_new(Arc::new(schema), vec![col_data]).unwrap(); + let expected = datatypes::arrow::array::StringArray::from(vec!["B", "C", "B", "C", "D"]); + + // Filter the batch + let filtered_batch = batch_filter(&batch, &physical_expr).unwrap(); + + // Expected: rows with col in ("B", "C", "D") + // That would be rows 0, 1, 3, 4, 5 + assert_eq!(filtered_batch.num_rows(), 5); + + let col_filtered = filtered_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col_filtered, &expected); + } + + #[test] + fn test_maybe_build_regex() { + // Test case for RegexMatch (case sensitive, non-negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(!negative); + assert!(regex.unwrap().is_match("axxb")); + + // Test case for RegexIMatch (case insensitive, non-negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexIMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(!negative); + assert!(regex.unwrap().is_match("AxxB")); + + // Test case for RegexNotMatch (case sensitive, negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexNotMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(negative); + + // Test case for RegexNotIMatch (case insensitive, negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexNotIMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(negative); + + // Test with empty regex pattern + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Utf8(Some("".to_string())), + ) + .unwrap(); + assert!(regex.is_none()); + assert!(!negative); + + // Test with non-regex operator + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::Eq, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_none()); + assert!(!negative); + + // Test with invalid regex pattern + let result = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Utf8(Some("a(b".to_string())), + ); + assert!(result.is_err()); + + // Test with non-string value + let result = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Int64(Some(123)), + ); + assert!(result.is_err()); + + // Test with null value + let result = + SimplePredicate::maybe_build_regex(Operator::RegexMatch, &ScalarValue::Utf8(None)); + assert!(result.is_err()); + } + + #[test] + fn test_regex_match_dictionary_array() { + use datatypes::arrow::array::StringDictionaryBuilder; + + // Create a StringDictionaryArray + let mut builder = StringDictionaryBuilder::::new(); + builder.append("apple").unwrap(); + builder.append("banana").unwrap(); + builder.append("apple").unwrap(); + builder.append("cherry").unwrap(); + let dict_array = builder.finish(); + + // Test regex that matches "apple" + let regex = regex::Regex::new(r"app.*").unwrap(); + let result = regexp_is_match_dictionary(&dict_array, Some(®ex)).unwrap(); + + // Should match indices 0 and 2 (both "apple") + assert_eq!(result.len(), 4); + assert!(result.value(0)); // "apple" + assert!(!result.value(1)); // "banana" + assert!(result.value(2)); // "apple" + assert!(!result.value(3)); // "cherry" + + // Test regex that matches "banana" + let regex2 = regex::Regex::new(r"ban.*").unwrap(); + let result2 = regexp_is_match_dictionary(&dict_array, Some(®ex2)).unwrap(); + + assert!(!result2.value(0)); // "apple" + assert!(result2.value(1)); // "banana" + assert!(!result2.value(2)); // "apple" + assert!(!result2.value(3)); // "cherry" + + // Test with no regex (should match all) + let result3 = regexp_is_match_dictionary(&dict_array, None).unwrap(); + assert!(result3.value(0)); + assert!(result3.value(1)); + assert!(result3.value(2)); + assert!(result3.value(3)); + } + + #[test] + fn test_regex_scan_masks_preserve_sql_null_semantics() { + let plain = Arc::new(datatypes::arrow::array::StringArray::from(vec![ + Some("api"), + Some("API"), + Some("db"), + None, + ])) as ArrayRef; + assert_regex_scan_masks( + &plain, + [ + vec![true, false, false, false], + vec![true, true, false, false], + vec![false, true, true, false], + vec![false, false, true, false], + ], + ); + + let dictionary = DictionaryArray::new( + datatypes::arrow::array::UInt32Array::from(vec![ + Some(0), + Some(1), + Some(2), + None, + Some(3), + ]), + Arc::new(datatypes::arrow::array::StringArray::from(vec![ + Some("api"), + Some("API"), + Some("db"), + None, + ])), + ); + let raw = + regexp_is_match_dictionary(&dictionary, Some(&Regex::new("^api$").unwrap())).unwrap(); + assert!(raw.is_null(3)); // null dictionary key + assert!(raw.is_null(4)); // non-null key referencing a null dictionary value + + let dictionary = Arc::new(dictionary) as ArrayRef; + assert_regex_scan_masks( + &dictionary, + [ + vec![true, false, false, false, false], + vec![true, true, false, false, false], + vec![false, true, true, false, false], + vec![false, false, true, false, false], + ], + ); + + let negative = regex_evaluator(Operator::RegexNotMatch); + assert!(!negative.evaluate_scalar(&ScalarValue::Utf8(None)).unwrap()); + } + + fn assert_regex_scan_masks(input: &ArrayRef, expected: [Vec; 4]) { + for (op, expected) in [ + Operator::RegexMatch, + Operator::RegexIMatch, + Operator::RegexNotMatch, + Operator::RegexNotIMatch, + ] + .into_iter() + .zip(expected) + { + assert_eq!( + BooleanBuffer::from(expected), + regex_evaluator(op).evaluate_array(input).unwrap(), + "{op:?}" + ); + } + } + + fn regex_evaluator(op: Operator) -> SimpleFilterEvaluator { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("host"))), + op, + right: Box::new("^api$".lit()), + }); + SimpleFilterEvaluator::try_new(&expr).unwrap() + } +} diff --git a/src/common/recordbatch/src/filter/tests.rs b/src/common/recordbatch/src/filter/tests.rs new file mode 100644 index 00000000000..77abf59b404 --- /dev/null +++ b/src/common/recordbatch/src/filter/tests.rs @@ -0,0 +1,72 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::execution::context::ExecutionProps; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion::logical_expr::{col, lit}; +use datafusion::physical_expr::create_physical_expr; +use datafusion_common::DFSchema; +use datatypes::arrow::datatypes::{DataType, Field, Schema}; + +use datatypes::arrow::array::{BooleanArray, RecordBatch}; + +use crate::filter::{batch_filter, boolean_array_to_scan_mask}; + +#[test] +fn batch_filter_test() { + let expr = col("ts").gt(lit(123456u64)); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("ts", DataType::UInt64, false), + ]); + let df_schema = DFSchema::try_from(schema.clone()).unwrap(); + let props = ExecutionProps::new(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(datatypes::arrow::array::Int32Array::from(vec![4, 5, 6])), + Arc::new(datatypes::arrow::array::UInt64Array::from(vec![ + 123456, 123457, 123458, + ])), + ], + ) + .unwrap(); + let new_batch = batch_filter(&batch, &physical_expr).unwrap(); + assert_eq!(new_batch.num_rows(), 2); + let first_column_values = new_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let expected = datatypes::arrow::array::Int32Array::from(vec![5, 6]); + assert_eq!(first_column_values, &expected); +} + +#[test] +fn test_nullable_boolean_predicate_becomes_scan_mask() { + let predicate = BooleanArray::from(vec![Some(true), None, Some(false)]); + assert_eq!( + BooleanArray::from(vec![true, false, false]), + boolean_array_to_scan_mask(&predicate) + ); +} diff --git a/src/common/recordbatch/src/filter/timestamp.rs b/src/common/recordbatch/src/filter/timestamp.rs new file mode 100644 index 00000000000..4c19aad546b --- /dev/null +++ b/src/common/recordbatch/src/filter/timestamp.rs @@ -0,0 +1,439 @@ +// 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. + +//! Timestamp-unit adaptation for predicates evaluated against historical SSTs. + +use common_time::timestamp::div_mod_units; +use datafusion::logical_expr::Operator; +use datafusion_common::ScalarValue; +use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar}; +use datatypes::arrow::datatypes::TimeUnit; + +use crate::filter::predicate::SimplePredicate; + +/// Builds a tz-naive timestamp literal for `value` in `unit`. +pub fn timestamp_scalar_value(value: i64, unit: TimeUnit) -> ScalarValue { + match unit { + TimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None), + TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None), + TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None), + TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None), + } +} + +impl SimplePredicate { + pub(crate) fn cast_timestamp_unit(&self, target_unit: TimeUnit) -> Option { + // Constants represent scan outcomes here. As in the public Matched/Pruned + // result, UNKNOWN can be collapsed to false under these AND/OR-only trees. + match self { + Self::And(left, right) | Self::Or(left, right) => { + let left = left.cast_timestamp_unit(target_unit)?; + let right = right.cast_timestamp_unit(target_unit)?; + let and = matches!(self, Self::And(..)); + Some(match (left, right) { + (Self::Constant(Some(false)), _) | (_, Self::Constant(Some(false))) if and => { + Self::Constant(Some(false)) + } + (Self::Constant(Some(true)), _) | (_, Self::Constant(Some(true))) if !and => { + Self::Constant(Some(true)) + } + (Self::Constant(Some(true)), other) | (other, Self::Constant(Some(true))) + if and => + { + other + } + (Self::Constant(Some(false)), other) | (other, Self::Constant(Some(false))) + if !and => + { + other + } + (left, right) => { + if and { + Self::And(Box::new(left), Box::new(right)) + } else { + Self::Or(Box::new(left), Box::new(right)) + } + } + }) + } + Self::IsNull { .. } => Some(self.clone()), + Self::Constant(value) => Some(Self::Constant(Some(value.unwrap_or(false)))), + Self::OrEqChain { literals } => { + let mut converted = Vec::with_capacity(literals.len()); + for literal in literals { + let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?; + let (value, unit) = timestamp_scalar_parts(&scalar)?; + let Some(value) = value else { + continue; + }; + let cast = div_mod_units(value, unit.into(), target_unit.into())?; + if cast.remainder == 0 { + converted.push(timestamp_scalar(cast.quotient, target_unit)?); + } + } + if converted.is_empty() { + return Some(Self::Constant(Some(false))); + } + Some(Self::OrEqChain { + literals: converted, + }) + } + Self::Comparison { literal, op, .. } => { + let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?; + let (value, unit) = timestamp_scalar_parts(&scalar)?; + let Some(value) = value else { + return Some(Self::Constant(Some(false))); + }; + if unit == target_unit { + return Some(self.clone()); + } + let cast = div_mod_units(value, unit.into(), target_unit.into())?; + let divisible = cast.remainder == 0; + let literal = timestamp_scalar(cast.quotient, target_unit)?; + let filter = |op| Self::Comparison { + literal: literal.clone(), + op, + regex: None, + regex_negative: false, + }; + Some(match op { + Operator::Eq if divisible => filter(Operator::Eq), + Operator::Eq => Self::Constant(Some(false)), + Operator::NotEq if divisible => filter(Operator::NotEq), + Operator::NotEq => Self::Constant(Some(true)), + // The quotient is floor(L), including negative timestamps. + // v > L is v > quotient whether or not L is representable. + Operator::Gt => filter(Operator::Gt), + Operator::GtEq if divisible => filter(Operator::GtEq), + // Between target values, v >= L is v > quotient. + Operator::GtEq => filter(Operator::Gt), + Operator::Lt if divisible => filter(Operator::Lt), + // Between target values, v < L is v <= quotient. + Operator::Lt => filter(Operator::LtEq), + Operator::LtEq => filter(Operator::LtEq), + _ => return None, + }) + } + } + } +} + +/// Extracts the value and unit from a tz-naive timestamp scalar. +fn timestamp_scalar_parts(scalar: &ScalarValue) -> Option<(Option, TimeUnit)> { + let (value, unit, timezone) = match scalar { + ScalarValue::TimestampSecond(v, tz) => (*v, TimeUnit::Second, tz), + ScalarValue::TimestampMillisecond(v, tz) => (*v, TimeUnit::Millisecond, tz), + ScalarValue::TimestampMicrosecond(v, tz) => (*v, TimeUnit::Microsecond, tz), + ScalarValue::TimestampNanosecond(v, tz) => (*v, TimeUnit::Nanosecond, tz), + _ => return None, + }; + // A timezone-aware literal doesn't compare against a tz-naive column; + // leave it to the caller instead of guessing the intended semantics. + (timezone.is_none()).then_some((value, unit)) +} + +/// Builds a one-element scalar array holding `value` in `unit`. +fn timestamp_scalar(value: i64, unit: TimeUnit) -> Option> { + timestamp_scalar_value(value, unit).to_scalar().ok() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::execution::context::ExecutionProps; + use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; + use datafusion::logical_expr::{Expr, col, lit}; + use datafusion::physical_expr::create_physical_expr; + use datafusion_common::arrow::array::ArrayRef; + use datafusion_common::arrow::buffer::BooleanBuffer; + use datafusion_common::cast::as_boolean_array; + use datafusion_common::{DFSchema, ScalarValue}; + use datatypes::arrow::array::{ + RecordBatch, TimestampMillisecondArray, TimestampNanosecondArray, + }; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::data_type::ConcreteDataType; + + use crate::filter::{SimpleFilterEvaluator, TimestampUnitCast, boolean_array_to_scan_mask}; + + #[test] + fn compound_timestamp_filters_preserve_unit_conversion() { + use datatypes::arrow::array::TimestampMicrosecondArray; + + let null = || lit(ScalarValue::TimestampMicrosecond(None, None)); + let predicates = [ + col("ts") + .gt(lit(ts_us(-1500))) + .and(col("ts").lt(lit(ts_us(500)))), + col("ts") + .eq(lit(ts_us(500))) + .or(col("ts").gt_eq(lit(ts_us(1500)))), + col("ts") + .is_not_null() + .and(col("ts").not_eq(lit(ts_us(500)))), + col("ts").eq(null()).or(col("ts").eq(lit(ts_us(-1000)))), + col("ts") + .eq(lit(ts_us(500))) + .and(col("ts").gt(lit(ts_us(0)))), + col("ts") + .not_eq(lit(ts_us(500))) + .or(col("ts").lt(lit(ts_us(0)))), + ]; + let native: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![ + -2000, -1000, 0, 1000, 2000, + ])); + let old: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![-2, -1, 0, 1, 2])); + let schema = Arc::new(Schema::new(vec![Field::new( + "ts", + native.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![native]).unwrap(); + for expr in predicates { + let physical = create_physical_expr( + &expr, + &DFSchema::try_from(schema.as_ref().clone()).unwrap(), + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); + let result = physical + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + let expected = boolean_array_to_scan_mask(as_boolean_array(&result).unwrap()) + .values() + .clone(); + let actual = match cast_to_ms(&expr).unwrap() { + TimestampUnitCast::Matched => BooleanBuffer::new_set(old.len()), + TimestampUnitCast::Pruned => BooleanBuffer::new_unset(old.len()), + TimestampUnitCast::Filter(filter) => filter.evaluate_array(&old).unwrap(), + }; + assert_eq!(actual, expected, "{expr}"); + } + } + + fn ts_us(v: i64) -> ScalarValue { + ScalarValue::TimestampMicrosecond(Some(v), None) + } + + fn cast_to_ms(expr: &Expr) -> Option { + SimpleFilterEvaluator::try_new(expr) + .unwrap() + .cast_timestamp_unit(&ConcreteDataType::timestamp_millisecond_datatype()) + } + + /// Evaluates a cast `Filter` outcome against millisecond values and + /// returns the mask. + fn eval_ms_mask(cast: Option, values: &[i64]) -> Vec { + let filter = match cast.expect("cast must apply") { + TimestampUnitCast::Filter(filter) => filter, + other => panic!("expected Filter outcome, got {other:?}"), + }; + let array = Arc::new(TimestampMillisecondArray::from(values.to_vec())) as ArrayRef; + filter.evaluate_array(&array).unwrap().iter().collect() + } + + #[test] + fn cast_timestamp_unit_converts_representable_literal() { + // ts = 7_000_000us evaluated against a ms column becomes ts = 7000ms. + assert_eq!( + vec![false, true, false], + eval_ms_mask( + cast_to_ms(&col("ts").eq(lit(ts_us(7_000_000)))), + &[6_999, 7_000, 7_001] + ) + ); + // != complements =. + assert_eq!( + vec![true, false, true], + eval_ms_mask( + cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_000)))), + &[6_999, 7_000, 7_001] + ) + ); + } + + #[test] + fn cast_timestamp_unit_prunes_non_representable_equality() { + // 7_000_500us is not a whole millisecond: no ms row can equal it, + // and every ms row satisfies `!=`. + assert!(matches!( + cast_to_ms(&col("ts").eq(lit(ts_us(7_000_500)))), + Some(TimestampUnitCast::Pruned) + )); + assert!(matches!( + cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_500)))), + Some(TimestampUnitCast::Matched) + )); + } + + #[test] + fn cast_timestamp_unit_strengthens_inequalities() { + // 2_500_500us is strictly between 2500ms and 2501ms: the cast must + // not round the literal to a boundary the original predicate + // excludes (>= 2500ms would wrongly match 2500ms). + let values = vec![2_499, 2_500, 2_501]; + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_500)))), &values) + ); + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_500)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_500)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_500)))), &values) + ); + + // Representable boundary keeps the original operator. + let values = vec![2_499, 2_500, 2_501]; + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_000)))), &values) + ); + assert_eq!( + vec![false, true, true], + eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_000)))), &values) + ); + assert_eq!( + vec![true, false, false], + eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_000)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_000)))), &values) + ); + } + + #[test] + fn cast_timestamp_unit_handles_negative_instants() { + // -2_500_500us floors to -2501ms with remainder 500us: only rows + // with instant >= -2.5005ms may match >= / >. + let values = vec![-2_502, -2_501, -2_500]; + assert_eq!( + vec![false, false, true], + eval_ms_mask( + cast_to_ms(&col("ts").gt_eq(lit(ts_us(-2_500_500)))), + &values + ) + ); + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(-2_500_500)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask( + cast_to_ms(&col("ts").lt_eq(lit(ts_us(-2_500_500)))), + &values + ) + ); + // Exactly representable negative literal: -2_500_000us == -2500ms. + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").eq(lit(ts_us(-2_500_000)))), &values) + ); + } + + #[test] + fn cast_timestamp_unit_or_chain_drops_unrepresentable_literals() { + // Only 7_000_000us is a whole millisecond; 7_000_500us cannot match. + let expr = col("ts") + .eq(lit(ts_us(7_000_500))) + .or(col("ts").eq(lit(ts_us(7_000_000)))); + assert_eq!( + vec![false, true], + eval_ms_mask(cast_to_ms(&expr), &[6_999, 7_000]) + ); + + // A chain where no literal is representable matches nothing. + let expr = col("ts") + .eq(lit(ts_us(7_000_500))) + .or(col("ts").eq(lit(ts_us(6_000_500)))); + assert!(matches!(cast_to_ms(&expr), Some(TimestampUnitCast::Pruned))); + } + + #[test] + fn cast_timestamp_unit_null_literal_prunes() { + // A NULL literal never compares true. + assert!(matches!( + cast_to_ms(&col("ts").eq(lit(ScalarValue::TimestampMicrosecond(None, None)))), + Some(TimestampUnitCast::Pruned) + )); + } + + #[test] + fn cast_timestamp_unit_same_unit_returns_filter_directly() { + // A literal already in the target unit is returned unchanged. + let expr = col("ts").eq(lit(ts_us(7_000_000))); + let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); + let cast = filter + .cast_timestamp_unit(&ConcreteDataType::timestamp_microsecond_datatype()) + .unwrap(); + let TimestampUnitCast::Filter(f) = cast else { + panic!("expected Filter, got {cast:?}") + }; + assert!(filter.is_eq() && f.is_eq()); + assert_eq!(filter.column_name(), f.column_name()); + assert_eq!(filter.literal_value(), f.literal_value()); + } + + #[test] + fn cast_timestamp_unit_rejects_non_timestamps() { + // Non-timestamp target type. + assert!( + SimpleFilterEvaluator::try_new(&col("ts").gt(lit(ts_us(1)))) + .unwrap() + .cast_timestamp_unit(&ConcreteDataType::int64_datatype()) + .is_none() + ); + // Non-timestamp literal against a timestamp target. + assert!(cast_to_ms(&col("ts").gt(lit(42_i64))).is_none()); + } + + #[test] + fn cast_timestamp_unit_to_finer_unit() { + // 5 seconds against a nanosecond column becomes 5_000_000_000ns. + let filter = SimpleFilterEvaluator::try_new( + &col("ts").eq(lit(ScalarValue::TimestampSecond(Some(5), None))), + ) + .unwrap() + .cast_timestamp_unit(&ConcreteDataType::timestamp_nanosecond_datatype()) + .and_then(|cast| match cast { + TimestampUnitCast::Filter(filter) => Some(filter), + _ => None, + }) + .unwrap(); + let array = Arc::new(TimestampNanosecondArray::from(vec![ + 4_999_999_999, + 5_000_000_000, + ])) as ArrayRef; + assert_eq!( + vec![false, true], + filter + .evaluate_array(&array) + .unwrap() + .iter() + .collect::>() + ); + } +} diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 4fb37cd5b11..0ccb8a4c28e 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(never_type)] - pub mod adapter; pub mod cursor; pub mod error; diff --git a/src/common/runtime/Cargo.toml b/src/common/runtime/Cargo.toml index 927aa17b9c2..410173f6445 100644 --- a/src/common/runtime/Cargo.toml +++ b/src/common/runtime/Cargo.toml @@ -25,13 +25,11 @@ futures.workspace = true lazy_static.workspace = true num_cpus.workspace = true once_cell.workspace = true -parking_lot.workspace = true paste.workspace = true pin-project.workspace = true prometheus.workspace = true ratelimit.workspace = true serde.workspace = true -serde_json.workspace = true snafu.workspace = true tempfile.workspace = true tokio.workspace = true diff --git a/src/common/sql/Cargo.toml b/src/common/sql/Cargo.toml index 3a66ec2d04e..318685996d5 100644 --- a/src/common/sql/Cargo.toml +++ b/src/common/sql/Cargo.toml @@ -10,7 +10,6 @@ common-base.workspace = true common-decimal.workspace = true common-error.workspace = true common-macro.workspace = true -common-telemetry.workspace = true common-time.workspace = true datafusion-sql.workspace = true datatypes.workspace = true diff --git a/src/common/stat/Cargo.toml b/src/common/stat/Cargo.toml index 2a99107c909..fcf927e9cca 100644 --- a/src/common/stat/Cargo.toml +++ b/src/common/stat/Cargo.toml @@ -8,7 +8,6 @@ license.workspace = true common-base.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -lazy_static.workspace = true num_cpus.workspace = true prometheus.workspace = true sysinfo.workspace = true diff --git a/src/common/substrait/Cargo.toml b/src/common/substrait/Cargo.toml index a4e78475051..a56b56b8301 100644 --- a/src/common/substrait/Cargo.toml +++ b/src/common/substrait/Cargo.toml @@ -13,7 +13,6 @@ bytes.workspace = true common-error.workspace = true common-function.workspace = true common-macro.workspace = true -common-telemetry.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true @@ -21,7 +20,3 @@ datafusion-substrait.workspace = true promql.workspace = true prost.workspace = true snafu.workspace = true - -[dev-dependencies] -datatypes.workspace = true -tokio.workspace = true diff --git a/src/common/substrait/src/df_substrait.rs b/src/common/substrait/src/df_substrait.rs index 08eeedd8316..9570076ee7a 100644 --- a/src/common/substrait/src/df_substrait.rs +++ b/src/common/substrait/src/df_substrait.rs @@ -25,7 +25,9 @@ use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::optimizer::AnalyzerRule; use datafusion::prelude::SessionConfig; use datafusion_expr::LogicalPlan; -use datafusion_substrait::logical_plan::consumer::from_substrait_plan; +use datafusion_substrait::logical_plan::consumer::{ + SubstraitConsumer, from_substrait_plan, from_substrait_plan_with_consumer, +}; use datafusion_substrait::logical_plan::producer::to_substrait_plan; use datafusion_substrait::substrait::proto::Plan; use prost::Message; @@ -74,6 +76,21 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor { } impl DFLogicalSubstraitConvertor { + /// Decode with a custom consumer, retaining Greptime's state-aggregate ordering correction. + pub async fn decode_with_consumer( + &self, + plan: &Plan, + state: &SessionState, + consumer: &impl SubstraitConsumer, + ) -> Result { + let df_plan = from_substrait_plan_with_consumer(consumer, plan) + .await + .context(DecodeDfPlanSnafu)?; + FixStateUdafOrderingAnalyzer {} + .analyze(df_plan, state.config_options()) + .context(DecodeDfPlanSnafu) + } + pub fn to_sub_plan( &self, plan: &LogicalPlan, diff --git a/src/common/substrait/src/lib.rs b/src/common/substrait/src/lib.rs index 3eeb3c570e1..a6932451a62 100644 --- a/src/common/substrait/src/lib.rs +++ b/src/common/substrait/src/lib.rs @@ -20,6 +20,7 @@ use async_trait::async_trait; use bytes::{Buf, Bytes}; use datafusion::execution::context::SessionState; pub use datafusion::execution::registry::SerializerRegistry; +pub use datafusion_substrait::extensions::Extensions; /// Re-export the Substrait module of datafusion, pub use datafusion_substrait::substrait as substrait_proto_df; pub use datafusion_substrait::{logical_plan as df_logical_plan, variation_const}; diff --git a/src/common/telemetry/src/lib.rs b/src/common/telemetry/src/lib.rs index 26bf5d53b3c..6daf2662abe 100644 --- a/src/common/telemetry/src/lib.rs +++ b/src/common/telemetry/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(duration_constructors)] - pub mod logging; mod macros; pub mod metric; diff --git a/src/common/telemetry/src/logging.rs b/src/common/telemetry/src/logging.rs index a9b50930786..dd99e901f75 100644 --- a/src/common/telemetry/src/logging.rs +++ b/src/common/telemetry/src/logging.rs @@ -348,7 +348,7 @@ impl Default for SlowQueryOptions { record_type: SlowQueriesRecordType::SystemTable, threshold: Duration::from_secs(30), sample_ratio: 1.0, - ttl: Duration::from_days(90), + ttl: Duration::from_secs(90 * 86400), } } } diff --git a/src/common/wal/src/config.rs b/src/common/wal/src/config.rs index 343df193ab4..0cbebfb20ab 100644 --- a/src/common/wal/src/config.rs +++ b/src/common/wal/src/config.rs @@ -138,7 +138,7 @@ mod tests { use super::*; use crate::TopicSelectorType; - use crate::config::object_store::CorruptedSegmentAction; + use crate::config::object_store::{AckMode, CorruptedSegmentAction}; use crate::config::{DatanodeKafkaConfig, MetasrvKafkaConfig}; #[test] @@ -234,7 +234,7 @@ mod tests { create_topic_timeout: Duration::from_secs(30), }, auto_create_topics: true, - auto_prune_interval: Duration::from_mins(30), + auto_prune_interval: Duration::from_secs(30 * 60), auto_prune_logical_delete: false, auto_prune_parallelism: 10, flush_trigger_size: ReadableSize::mb(512), @@ -294,6 +294,9 @@ mod tests { assert_eq!(config.prefix, "wal"); assert_eq!(config.flush_interval, Duration::from_millis(100)); assert_eq!(config.max_batch_bytes, ReadableSize::mb(8)); + assert_eq!(config.ack_mode, AckMode::Durable); + assert_eq!(config.max_unpersisted_bytes, ReadableSize::mb(64)); + assert_eq!(config.max_unpersisted_age, Duration::from_secs(8)); assert_eq!(config.on_corrupted_segment, CorruptedSegmentAction::Skip); assert!(matches!( MetasrvWalConfig::try_from(datanode_wal_config).unwrap_err(), @@ -306,6 +309,9 @@ mod tests { prefix = "cluster-a/wal" flush_interval = "500ms" max_batch_bytes = "4MB" + ack_mode = "enqueued" + max_unpersisted_bytes = "32MB" + max_unpersisted_age = "4s" on_corrupted_segment = "fail" "#; let datanode_wal_config: DatanodeWalConfig = toml::from_str(toml_str).unwrap(); @@ -314,12 +320,24 @@ mod tests { prefix: "cluster-a/wal".to_string(), flush_interval: Duration::from_millis(500), max_batch_bytes: ReadableSize::mb(4), + ack_mode: AckMode::Enqueued, + max_unpersisted_bytes: ReadableSize::mb(32), + max_unpersisted_age: Duration::from_secs(4), on_corrupted_segment: CorruptedSegmentAction::Fail, }; assert_eq!( datanode_wal_config, DatanodeWalConfig::ObjectStore(expected) ); + let serialized = toml::to_string(&datanode_wal_config).unwrap(); + let table: toml::Table = toml::from_str(&serialized).unwrap(); + assert_eq!(table["ack_mode"].as_str(), Some("enqueued")); + assert_eq!(table["max_unpersisted_bytes"].as_str(), Some("32MiB")); + assert_eq!(table["max_unpersisted_age"].as_str(), Some("4s")); + assert_eq!( + datanode_wal_config, + toml::from_str::(&serialized).unwrap() + ); // The persisted tag is not accepted as a config provider. let toml_str = r#" diff --git a/src/common/wal/src/config/kafka/common.rs b/src/common/wal/src/config/kafka/common.rs index cf632fbdce3..daa5607e476 100644 --- a/src/common/wal/src/config/kafka/common.rs +++ b/src/common/wal/src/config/kafka/common.rs @@ -38,7 +38,7 @@ pub const DEFAULT_BACKOFF_CONFIG: BackoffConfig = BackoffConfig { }; /// Default interval for auto WAL pruning. -pub const DEFAULT_AUTO_PRUNE_INTERVAL: Duration = Duration::from_mins(30); +pub const DEFAULT_AUTO_PRUNE_INTERVAL: Duration = Duration::from_secs(30 * 60); /// Default mode for auto WAL pruning. pub const DEFAULT_AUTO_PRUNE_LOGICAL_DELETE: bool = false; /// Default limit for concurrent auto pruning tasks. diff --git a/src/common/wal/src/config/object_store.rs b/src/common/wal/src/config/object_store.rs index 04a212fb53f..dd2c4544129 100644 --- a/src/common/wal/src/config/object_store.rs +++ b/src/common/wal/src/config/object_store.rs @@ -24,6 +24,18 @@ use serde::{Deserialize, Serialize}; /// node id and the generation from the metasrv. pub const STANDALONE_GENERATION: u64 = 0; +/// When an append to the object store WAL returns. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum AckMode { + /// An append returns after the object holding its entries is durable. + #[default] + Durable, + /// An append returns once its entries are admitted and their ids are + /// assigned; the object is created in the background. + Enqueued, +} + /// What a read does with a segment that still does not decode after a /// second fetch, because its checksum does not match or its content /// disagrees with its footer entry. @@ -56,6 +68,15 @@ pub struct ObjectStoreWalConfig { pub flush_interval: Duration, /// The max size of a single batch object, defaults to 8MiB. pub max_batch_bytes: ReadableSize, + /// When an append returns, defaults to `durable`. + pub ack_mode: AckMode, + /// Size of the unpersisted backlog at which appends stall until an + /// upload completes in `enqueued` mode, defaults to 64MiB. + pub max_unpersisted_bytes: ReadableSize, + /// Age of the oldest unpersisted entry at which appends stall until an + /// upload completes in `enqueued` mode, defaults to 8s. + #[serde(with = "humantime_serde")] + pub max_unpersisted_age: Duration, /// What a read does with a segment that still does not decode after a /// second fetch, because its checksum does not match or its content /// disagrees with its footer entry, defaults to `skip`. @@ -69,6 +90,9 @@ impl Default for ObjectStoreWalConfig { prefix: "wal".to_string(), flush_interval: Duration::from_millis(100), max_batch_bytes: ReadableSize::mb(8), + ack_mode: AckMode::Durable, + max_unpersisted_bytes: ReadableSize::mb(64), + max_unpersisted_age: Duration::from_secs(8), on_corrupted_segment: CorruptedSegmentAction::Skip, } } diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 6b78858bd2f..4d4e470e665 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -16,6 +16,7 @@ api.workspace = true arrow-flight.workspace = true async-trait.workspace = true bytes.workspace = true +catalog.workspace = true client.workspace = true common-base.workspace = true common-catalog.workspace = true @@ -56,6 +57,7 @@ metric-engine.workspace = true mito2.workspace = true num_cpus.workspace = true object-store.workspace = true +partition.workspace = true prometheus.workspace = true prost.workspace = true query.workspace = true @@ -75,11 +77,9 @@ tracing.workspace = true [dev-dependencies] cache.workspace = true -client.workspace = true common-meta = { workspace = true, features = ["testing"] } common-query.workspace = true common-test-util.workspace = true datafusion-common.workspace = true mito2 = { workspace = true, features = ["test"] } -partition.workspace = true session.workspace = true diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index ae4b0723e1a..cfbcf948b96 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -19,26 +19,32 @@ use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; +use catalog::information_schema::NoopInformationExtension; +use catalog::kvbackend::KvBackendCatalogManagerBuilder; +use client::client_manager::NodeClients; use common_base::Plugins; use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE}; use common_datasource::object_store::LocalFileAccess; use common_error::ext::BoxedError; use common_greptimedb_telemetry::GreptimeDBTelemetryTask; +use common_grpc::channel_manager::ChannelConfig; use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::datanode::TopicStatsReporter; use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef}; use common_meta::kv_backend::KvBackendRef; +use common_meta::node_manager::NodeManagerRef; pub use common_procedure::options::ProcedureConfig; use common_query::prelude::set_default_prefix; use common_stat::ResourceStatImpl; use common_telemetry::{error, info, warn}; use common_wal::config::DatanodeWalConfig; use common_wal::config::kafka::DatanodeKafkaConfig; -use common_wal::config::object_store::ObjectStoreWalConfig; +use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; use common_wal::config::raft_engine::RaftEngineConfig; use file_engine::engine::FileRegionEngine; +use log_store::ObjectStoreLogStore; use log_store::kafka::log_store::KafkaLogStore; use log_store::kafka::{GlobalIndexCollector, default_index_file}; use log_store::noop::log_store::NoopLogStore; @@ -49,12 +55,15 @@ use mito2::config::MitoConfig; use mito2::engine::{MitoEngine, MitoEngineBuilder}; use mito2::region::opener::PartitionExprFetcherRef; use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef}; +use object_store::ObjectStore; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::util::normalize_dir; use query::QueryEngineFactory; -use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef}; +use query::dummy_catalog::TableProviderFactoryRef; +use query::region_query::RegionQueryHandlerFactoryRef; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; +use store_api::logstore::LogStore; use store_api::path_utils::WAL_DIR; use store_api::region_engine::{ RegionEngineRef, RegionRole, SetRegionRoleStateResponse, SettableRegionRoleState, @@ -67,7 +76,7 @@ use crate::error::{ self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu, DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu, InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu, - ObjectStoreWalNotSupportedSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, + ObjectStoreWalNotStandaloneSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu, }; use crate::event_listener::{ @@ -77,6 +86,7 @@ use crate::event_listener::{ use crate::greptimedb_telemetry::get_greptimedb_telemetry_task; use crate::heartbeat::HeartbeatTask; use crate::partition_expr_fetcher::MetaPartitionExprFetcher; +use crate::region_query::DatanodeRegionQueryHandler; use crate::region_server::{DummyTableProviderFactory, RegionServer}; use crate::store::{self, new_object_store_without_cache}; use crate::utils::{RegionOpenRequests, build_region_open_requests}; @@ -90,6 +100,7 @@ pub struct Datanode { greptimedb_telemetry_task: Arc, leases_notifier: Option>, plugins: Plugins, + object_store_log_store: Option>, } impl Datanode { @@ -135,27 +146,53 @@ impl Datanode { self.services = services; } + /// Shuts down every component in order. A failed step does not skip the + /// later ones; the first error is returned and the rest are logged. pub async fn shutdown(&mut self) -> Result<()> { - self.services - .shutdown_all() - .await - .context(ShutdownServerSnafu)?; + let mut first_error = None; + record_shutdown_error( + &mut first_error, + self.services + .shutdown_all() + .await + .context(ShutdownServerSnafu), + ); let _ = self.greptimedb_telemetry_task.stop().await; if let Some(heartbeat_task) = &self.heartbeat_task { - heartbeat_task - .close() - .map_err(BoxedError::new) - .context(ShutdownInstanceSnafu)?; + record_shutdown_error( + &mut first_error, + heartbeat_task + .close() + .map_err(BoxedError::new) + .context(ShutdownInstanceSnafu), + ); } - self.region_server.stop().await?; - Ok(()) + record_shutdown_error(&mut first_error, self.region_server.stop().await); + if let Some(log_store) = &self.object_store_log_store { + record_shutdown_error( + &mut first_error, + log_store + .stop() + .await + .map_err(BoxedError::new) + .context(ShutdownInstanceSnafu), + ); + } + first_error.map_or(Ok(()), Err) } pub fn region_server(&self) -> RegionServer { self.region_server.clone() } + /// Returns the object store log store when the datanode runs on the object + /// store WAL. + #[cfg(any(test, feature = "testing"))] + pub fn object_store_log_store(&self) -> Option> { + self.object_store_log_store.clone() + } + pub fn plugins(&self) -> Plugins { self.plugins.clone() } @@ -173,6 +210,8 @@ pub struct DatanodeBuilder { local_file_access: LocalFileAccess, #[cfg(feature = "enterprise")] extension_range_provider_factory: Option, + /// Set as soon as the object store log store exists so that a failed build can stop it. + object_store_log_store: Option>, } impl DatanodeBuilder { @@ -189,6 +228,7 @@ impl DatanodeBuilder { #[cfg(feature = "enterprise")] extension_range_provider_factory: None, topic_stats_reporter: None, + object_store_log_store: None, } } @@ -256,6 +296,22 @@ impl DatanodeBuilder { } pub async fn build(mut self) -> Result { + self.try_build().await + } + + /// Builds the datanode and stops the object store log store if a later step fails. + async fn try_build(&mut self) -> Result { + let result = self.build_inner().await; + if result.is_err() + && let Some(log_store) = &self.object_store_log_store + && let Err(err) = log_store.stop().await + { + warn!(err; "Failed to stop the object store log store after a failed build"); + } + result + } + + async fn build_inner(&mut self) -> Result { validate_region_engine_config(&self.opts.region_engine)?; let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?; set_default_prefix(self.opts.default_column_prefix.as_deref()) @@ -269,6 +325,7 @@ impl DatanodeBuilder { // writable upon open. let controlled_by_metasrv = meta_client.is_some(); if let DatanodeWalConfig::ObjectStore(config) = &self.opts.wal { + ensure!(!controlled_by_metasrv, ObjectStoreWalNotStandaloneSnafu); validate_object_store_wal_config(config)?; } @@ -307,6 +364,7 @@ impl DatanodeBuilder { schema_metadata_manager, region_event_listener, file_ref_manager, + cache_registry.clone(), ) .await?; @@ -372,6 +430,7 @@ impl DatanodeBuilder { region_event_receiver, leases_notifier, plugins: self.plugins.clone(), + object_store_log_store: self.object_store_log_store.take(), }) } @@ -446,14 +505,41 @@ impl DatanodeBuilder { schema_metadata_manager: SchemaMetadataManagerRef, event_listener: RegionServerEventListenerRef, file_ref_manager: FileReferenceManagerRef, + cache_registry: Arc, ) -> Result { let opts: &DatanodeOptions = &self.opts; + // Keep `with_dist_planner` false: datanodes execute received MergeScan plans. + let catalog_manager = KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + self.kv_backend.clone(), + cache_registry.clone(), + ) + .build(); + let partition_manager = catalog_manager.partition_manager(); + + // MergeScan region requests stream until completion, so they have no client timeout. + let mut channel_config = ChannelConfig { + timeout: None, + ..Default::default() + }; + if opts.grpc.flight_compression.transport_compression() { + channel_config.accept_compression = true; + channel_config.send_compression = true; + } + let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config)); + let region_query_handler = + if let Some(factory) = self.plugins.get::() { + factory.build(partition_manager.clone(), node_manager) + } else { + DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager) + }; + let query_engine_factory = QueryEngineFactory::try_new_with_plugins( // query engine in datanode only executes plan with resolved table source. - DummyCatalogManager::arc(), - None, - None, + catalog_manager, + Some(partition_manager), + Some(region_query_handler), None, None, None, @@ -660,11 +746,57 @@ impl DatanodeBuilder { builder.try_build().await.context(BuildMitoEngineSnafu)? } - DatanodeWalConfig::ObjectStore(_) => return ObjectStoreWalNotSupportedSnafu.fail(), + DatanodeWalConfig::ObjectStore(object_store_config) => { + let object_store = wal_object_store( + &object_store_manager, + &opts.storage, + &object_store_config.storage_provider, + )?; + let node_id = opts.node_id.context(MissingNodeIdSnafu)?; + let log_store = + Self::build_object_store_log_store(object_store, object_store_config, node_id) + .await?; + self.object_store_log_store = Some(log_store.clone()); + + let builder = MitoEngineBuilder::new( + &opts.storage.data_home, + config, + log_store, + object_store_manager, + schema_metadata_manager, + file_ref_manager, + partition_expr_fetcher, + plugins, + ); + + #[cfg(feature = "enterprise")] + let builder = builder.with_extension_range_provider_factory( + self.extension_range_provider_factory.take(), + ); + + builder.try_build().await.context(BuildMitoEngineSnafu)? + } }; Ok(mito_engine) } + /// Builds [ObjectStoreLogStore] under the node prefix of `node_id` in the + /// standalone generation, the prefix the regions persist in their WAL options. + async fn build_object_store_log_store( + object_store: ObjectStore, + config: &ObjectStoreWalConfig, + node_id: u64, + ) -> Result> { + info!( + "Creating object store logstore with config: {:?}, node id: {}", + config, node_id + ); + ObjectStoreLogStore::try_new(object_store, config, node_id, STANDALONE_GENERATION) + .await + .map_err(Box::new) + .context(OpenLogStoreSnafu) + } + /// Builds [RaftEngineLogStore]. async fn build_raft_engine_log_store( data_home: &str, @@ -714,6 +846,43 @@ impl DatanodeBuilder { } } +/// Keeps the first shutdown error and logs the later ones. +pub(crate) fn record_shutdown_error(first_error: &mut Option, result: Result<()>) { + if let Err(err) = result { + if first_error.is_none() { + *first_error = Some(err); + } else { + warn!(err; "Ignored a later shutdown error"); + } + } +} + +/// Returns the object store that holds the WAL objects: the default one for an +/// empty `storage_provider`, otherwise the configured provider with that name. +fn wal_object_store( + object_store_manager: &ObjectStoreManager, + storage: &StorageConfig, + storage_provider: &str, +) -> Result { + if storage_provider.is_empty() { + return Ok(object_store_manager.default_object_store().clone()); + } + object_store_manager + .find(storage_provider) + .cloned() + .with_context(|| { + let configured = std::iter::once(&storage.store) + .chain(storage.providers.iter()) + .map(|store| store.config_name()) + .collect::>(); + InvalidObjectStoreWalConfigSnafu { + field: "storage_provider", + value: storage_provider, + reason: format!("must name a configured storage provider, one of {configured:?}"), + } + }) +} + /// Rejects repeated engine types before their configs can silently overwrite each other. fn validate_region_engine_config(configs: &[RegionEngineConfig]) -> Result<()> { let mut engine_indices = HashMap::new(); @@ -780,6 +949,22 @@ fn validate_object_store_wal_config(config: &ObjectStoreWalConfig) -> Result<()> reason: "must be greater than 0", } ); + ensure!( + config.max_unpersisted_bytes.as_bytes() > 0, + InvalidObjectStoreWalConfigSnafu { + field: "max_unpersisted_bytes", + value: config.max_unpersisted_bytes.to_string(), + reason: "must be greater than 0", + } + ); + ensure!( + config.max_unpersisted_age > Duration::ZERO, + InvalidObjectStoreWalConfigSnafu { + field: "max_unpersisted_age", + value: format!("{:?}", config.max_unpersisted_age), + reason: "must be greater than 0", + } + ); Ok(()) } @@ -907,31 +1092,42 @@ mod tests { use std::io::Write; use std::path::Path; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; - use cache::build_datanode_cache_registry; + use cache::build_datanode_layered_cache_registry; use common_base::Plugins; use common_base::readable_size::ReadableSize; use common_config::Configurable; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; - use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::key::RegionRoleSet; use common_meta::key::datanode_table::DatanodeTableManager; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::node_manager::NodeManagerRef; use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; use common_wal::config::DatanodeWalConfig; use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; + use log_store::ObjectStoreLogStore; + use log_store::error::Error as LogStoreError; + use meta_client::client::MetaClientBuilder; use mito2::engine::MITO_ENGINE_NAME; + use object_store::ObjectStore; + use object_store::manager::ObjectStoreManager; + use partition::manager::PartitionRuleManagerRef; + use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef}; + use store_api::logstore::LogStore; use store_api::region_request::RegionRequest; use store_api::storage::RegionId; - use crate::config::{DatanodeOptions, RegionEngineConfig}; + use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use crate::datanode::{ DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config, + wal_object_store, }; - use crate::error::Error; + use crate::error::{self, Error}; + use crate::region_query::DatanodeRegionQueryHandler; use crate::tests::{MockRegionEngine, mock_region_server}; async fn setup_table_datanode(kv: &KvBackendRef) { @@ -960,11 +1156,8 @@ mod tests { mock_region_server.register_engine(mock_region.clone()); let kv_backend = Arc::new(MemoryKvBackend::new()); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut builder = DatanodeBuilder::new( DatanodeOptions { @@ -1006,11 +1199,8 @@ mod tests { wal: DatanodeWalConfig, kv_backend: KvBackendRef, ) -> DatanodeBuilder { - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut opts = DatanodeOptions { node_id: Some(0), wal, @@ -1043,6 +1233,14 @@ mod tests { std::fs::read_dir(dir).unwrap().next().is_none() } + /// Returns true if the store rejects even an empty append, which it does only once stopped. + async fn is_stopped(log_store: &ObjectStoreLogStore) -> bool { + matches!( + log_store.append_batch(vec![]).await, + Err(LogStoreError::ObjectStoreWalStopped { .. }) + ) + } + #[test] fn test_validate_unique_region_engine_config() { let mut opts = DatanodeOptions::default(); @@ -1163,12 +1361,28 @@ mod tests { }, "max_batch_bytes", ), + ( + ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(0), + ..Default::default() + }, + "max_unpersisted_bytes", + ), + ( + ObjectStoreWalConfig { + max_unpersisted_age: Duration::ZERO, + ..Default::default() + }, + "max_unpersisted_age", + ), ]; for (config, expected_field) in cases { let expected_value = match expected_field { "prefix" => config.prefix.clone(), "flush_interval" => format!("{:?}", config.flush_interval), + "max_unpersisted_bytes" => config.max_unpersisted_bytes.to_string(), + "max_unpersisted_age" => format!("{:?}", config.max_unpersisted_age), _ => config.max_batch_bytes.to_string(), }; let data_home = create_temp_dir("object-store-wal-invalid-config"); @@ -1187,18 +1401,217 @@ mod tests { } #[tokio::test] - async fn test_build_rejects_object_store_wal_as_not_supported() { - let data_home = create_temp_dir("object-store-wal-not-supported"); + async fn test_build_rejects_object_store_wal_with_meta_client() { + let data_home = create_temp_dir("object-store-wal-meta-client"); + let meta_client = Arc::new(MetaClientBuilder::datanode_default_options(0).build()); + + let mut builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig::default(), + ); + builder.with_meta_client(meta_client); + let err = build_err(builder).await; + + assert_matches!(err, Error::ObjectStoreWalNotStandalone { .. }); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + // The builder stops before it creates any storage or log store. + assert!(is_empty_dir(data_home.path())); + } + + #[tokio::test] + async fn test_build_rejects_unknown_wal_storage_provider() { + let data_home = create_temp_dir("object-store-wal-unknown-provider"); + let builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig { + storage_provider: "s3".to_string(), + ..Default::default() + }, + ); + let err = build_err(builder).await; + + let Error::InvalidObjectStoreWalConfig { + field, + value, + reason, + .. + } = &err + else { + panic!("unexpected error: {err:?}"); + }; + assert_eq!(*field, "storage_provider"); + assert_eq!(value, "s3"); + assert!(reason.contains(r#"["File"]"#), "{reason}"); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + // No log store is created under the prefix. + assert!(!data_home.path().join("wal").exists()); + } + + #[test] + fn test_wal_object_store_selects_provider() { + fn fs_object_store(dir: &Path) -> ObjectStore { + let builder = object_store::services::Fs::default().root(dir.to_str().unwrap()); + ObjectStore::new(builder).unwrap() + } + + let default_dir = create_temp_dir("wal-object-store-default"); + let named_dir = create_temp_dir("wal-object-store-named"); + let mut manager = ObjectStoreManager::new("File", fs_object_store(default_dir.path())); + manager.add("Wal-Store", fs_object_store(named_dir.path())); + let storage = StorageConfig::default(); + + let root = |store: ObjectStore| store.info().root().clone(); + assert_eq!( + root(wal_object_store(&manager, &storage, "").unwrap()), + root(manager.default_object_store().clone()) + ); + assert_eq!( + root(wal_object_store(&manager, &storage, "wal-store").unwrap()), + root(manager.find("Wal-Store").unwrap().clone()) + ); + assert_matches!( + wal_object_store(&manager, &storage, "missing").unwrap_err(), + Error::InvalidObjectStoreWalConfig { + field: "storage_provider", + .. + } + ); + } + + struct CountingRegionQueryHandlerFactory(Arc); + + impl RegionQueryHandlerFactory for CountingRegionQueryHandlerFactory { + fn build( + &self, + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + ) -> query::region_query::RegionQueryHandlerRef { + self.0.fetch_add(1, Ordering::SeqCst); + DatanodeRegionQueryHandler::arc(partition_manager, node_manager) + } + } + + #[tokio::test] + async fn test_build_uses_installed_region_query_handler_factory() { + let data_home = create_temp_dir("region-query-handler-factory"); + let builds = Arc::new(AtomicUsize::new(0)); + let plugins = Plugins::new(); + plugins.insert(Arc::new(CountingRegionQueryHandlerFactory(builds.clone())) + as RegionQueryHandlerFactoryRef); + let mut builder = datanode_builder( + data_home.path().to_str().unwrap(), + DatanodeWalConfig::default(), + Arc::new(MemoryKvBackend::new()), + ); + builder.set_plugins(plugins); + + let mut datanode = builder.build().await.unwrap(); + assert_eq!(builds.load(Ordering::SeqCst), 1); + datanode.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn test_build_and_shutdown_with_raft_engine_wal() { + let data_home = create_temp_dir("raft-engine-wal"); + let builder = datanode_builder( + data_home.path().to_str().unwrap(), + DatanodeWalConfig::default(), + Arc::new(MemoryKvBackend::new()), + ); + + let mut datanode = builder.build().await.unwrap(); + assert!(datanode.object_store_log_store.is_none()); + assert!(data_home.path().join("wal").exists()); + datanode.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn test_build_and_shutdown_with_object_store_wal() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal"); let builder = object_store_wal_builder( data_home.path().to_str().unwrap(), ObjectStoreWalConfig::default(), ); - let err = build_err(builder).await; - assert_matches!(err, Error::ObjectStoreWalNotSupported { .. }); - assert_eq!(StatusCode::Unsupported, err.status_code()); - // No log store is created under the prefix. - assert!(!data_home.path().join("wal").exists()); + let mut datanode = builder.build().await.unwrap(); + let log_store = datanode.object_store_log_store.clone().unwrap(); + assert!(!is_stopped(&log_store).await); + // The store runs under the node prefix derived from the configured root. + assert!(format!("{log_store:?}").contains(r#"prefix: "wal/datanodes/0/epochs/0""#)); + + datanode.shutdown().await.unwrap(); + assert!(is_stopped(&log_store).await); + } + + #[tokio::test] + async fn test_build_with_named_wal_storage_provider() { + let data_home = create_temp_dir("object-store-wal-named-provider"); + // The default file store is registered under its provider name. + let builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig { + storage_provider: "file".to_string(), + ..Default::default() + }, + ); + + let mut datanode = builder.build().await.unwrap(); + assert!(datanode.object_store_log_store.is_some()); + datanode.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn test_shutdown_stops_every_engine_and_the_log_store_after_engine_failures() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal-shutdown-failure"); + let builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig::default(), + ); + let mut datanode = builder.build().await.unwrap(); + let log_store = datanode.object_store_log_store.clone().unwrap(); + // Both engines fail to stop, so whichever the region server stops first fails. + let stop_calls = Arc::new(AtomicUsize::new(0)); + for name in ["failing-a", "failing-b"] { + let stop_calls = stop_calls.clone(); + let (engine, _) = MockRegionEngine::with_custom_apply_fn(name, move |engine| { + engine.handle_stop_mock_fn = Some(Box::new(move || { + stop_calls.fetch_add(1, Ordering::Relaxed); + error::UnexpectedSnafu { + violated: "stop failed", + } + .fail() + })); + }); + datanode.region_server().register_engine(engine); + } + + let err = datanode.shutdown().await.unwrap_err(); + + assert_matches!(err, Error::StopRegionEngine { .. }); + assert_eq!(2, stop_calls.load(Ordering::Relaxed)); + assert!(is_stopped(&log_store).await); + } + + #[tokio::test] + async fn test_failed_build_stops_object_store_log_store() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal-failed-build"); + let kv_backend = Arc::new(MemoryKvBackend::new()); + // Regions registered in the metadata but missing on disk fail the region open + // that runs after the log store and the engines are built. + setup_table_datanode(&(kv_backend.clone() as _)).await; + let mut builder = datanode_builder( + data_home.path().to_str().unwrap(), + DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default()), + kv_backend, + ); + + assert!(builder.try_build().await.is_err()); + + let log_store = builder.object_store_log_store.as_ref().unwrap(); + assert!(is_stopped(log_store).await); } #[test] diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index b62708aa343..e9b5d5ec1dc 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -145,6 +145,13 @@ pub enum Error { source: Box, }, + #[snafu(display("Invalid query context: {}", reason))] + InvalidQueryContext { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid SQL, error: {}", msg))] InvalidSql { msg: String }, @@ -319,8 +326,8 @@ pub enum Error { location: Location, }, - #[snafu(display("Object store WAL is not supported yet"))] - ObjectStoreWalNotSupported { + #[snafu(display("Object store WAL is only supported in standalone mode"))] + ObjectStoreWalNotStandalone { #[snafu(implicit)] location: Location, }, @@ -485,7 +492,8 @@ impl ErrorExt for Error { Delete { source, .. } => source.status_code(), - InvalidSql { .. } + InvalidQueryContext { .. } + | InvalidSql { .. } | IllegalPrimaryKeysDef { .. } | MissingTimestampColumn { .. } | SchemaNotFound { .. } @@ -499,6 +507,7 @@ impl ErrorExt for Error { | TomlFormat { .. } | DuplicateRegionEngineConfig { .. } | InvalidObjectStoreWalConfig { .. } + | ObjectStoreWalNotStandalone { .. } | BuildDatanode { .. } => StatusCode::InvalidArguments, PayloadNotExist { .. } @@ -523,9 +532,7 @@ impl ErrorExt for Error { OpenLogStore { source, .. } => source.status_code(), MetaClientInit { source, .. } => source.status_code(), - UnsupportedOutput { .. } - | NotYetImplemented { .. } - | ObjectStoreWalNotSupported { .. } => StatusCode::Unsupported, + UnsupportedOutput { .. } | NotYetImplemented { .. } => StatusCode::Unsupported, HandleRegionRequest { source, .. } | GetRegionMetadata { source, .. } | HandleBatchOpenRequest { source, .. } diff --git a/src/datanode/src/lib.rs b/src/datanode/src/lib.rs index c14ec886fac..5a8948e2c84 100644 --- a/src/datanode/src/lib.rs +++ b/src/datanode/src/lib.rs @@ -22,6 +22,7 @@ pub mod heartbeat; pub mod metrics; mod partition_expr_fetcher; pub mod query_stream; +pub mod region_query; pub mod region_server; pub mod service; pub mod store; diff --git a/src/datanode/src/metrics.rs b/src/datanode/src/metrics.rs index 7573490c747..e9f947ef9da 100644 --- a/src/datanode/src/metrics.rs +++ b/src/datanode/src/metrics.rs @@ -85,6 +85,16 @@ lazy_static! { ) .unwrap(); + /// Failed region open attempts, including individual and batch requests. + /// Counts every request in an engine group when no individual results are available. + /// Failed retries count separately; explicitly ignored missing regions do not count. + pub static ref REGION_OPEN_FAILURES_TOTAL: IntCounterVec = register_int_counter_vec!( + "greptime_datanode_region_open_failures_total", + "Failed region open attempts, including retries and all requests in failed engine groups, excluding ignored missing regions", + &["engine", "status_code"] + ) + .unwrap(); + /// Total count of failed insert requests to region server, labeled with request /// type and status code. pub static ref REGION_SERVER_INSERT_FAIL_COUNT: IntCounterVec = register_int_counter_vec!( diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs new file mode 100644 index 00000000000..1642362ee1a --- /dev/null +++ b/src/datanode/src/region_query.rs @@ -0,0 +1,109 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The datanode's [`RegionQueryHandler`] implementation for querying remote regions in a +//! `MergeScan` plan. + +use std::sync::Arc; + +use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; +use async_trait::async_trait; +use common_error::ext::BoxedError; +use common_meta::node_manager::NodeManagerRef; +use common_query::request::QueryRequest; +use common_recordbatch::SendableRecordBatchStream; +use partition::manager::PartitionRuleManagerRef; +use query::error::{RegionQuerySnafu, Result as QueryResult, UnimplementedSnafu}; +use query::region_query::{RegionQueryHandler, RegionQueryTarget}; +use session::ReadPreference; +use snafu::ResultExt; +use store_api::storage::RegionId; + +/// Serves region queries by resolving their leader and forwarding requests to its datanode. +pub(crate) struct DatanodeRegionQueryHandler { + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, +} + +impl DatanodeRegionQueryHandler { + pub fn arc( + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + ) -> Arc { + Arc::new(Self { + partition_manager, + node_manager, + }) + } +} + +#[async_trait] +impl RegionQueryHandler for DatanodeRegionQueryHandler { + async fn select_target( + &self, + _read_preference: ReadPreference, + region_id: RegionId, + ) -> QueryResult { + // The leader of a region is the only peer that can serve a region query. + // + let peer = self + .partition_manager + .find_region_leader(region_id) + .await + .map_err(BoxedError::new) + .context(RegionQuerySnafu)?; + + Ok(RegionQueryTarget::new(peer)) + } + + async fn do_get( + &self, + target: &RegionQueryTarget, + request: QueryRequest, + ) -> QueryResult { + self.node_manager + .datanode(target.peer()) + .await + .handle_query(request) + .await + .map_err(BoxedError::new) + .context(RegionQuerySnafu) + } + + /// Datanodes forward received region queries and do not produce remote dynamic filters. + async fn handle_remote_dyn_filter_update( + &self, + _target: &RegionQueryTarget, + _query_id: String, + _update: RemoteDynFilterUpdate, + ) -> QueryResult<()> { + UnimplementedSnafu { + operation: "remote dynamic filter update from a datanode", + } + .fail() + } + + /// See [`RegionQueryHandler::handle_remote_dyn_filter_update`]. + async fn handle_remote_dyn_filter_unregister( + &self, + _target: &RegionQueryTarget, + _query_id: String, + _unregister: RemoteDynFilterUnregister, + ) -> QueryResult<()> { + UnimplementedSnafu { + operation: "remote dynamic filter unregister from a datanode", + } + .fail() + } +} diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 597fdd66d96..f8261f03e11 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -92,6 +92,7 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio::time::{self, timeout}; use tonic::{Request, Response, Result as TonicResult}; +use crate::datanode::record_shutdown_error; use crate::error::{ self, BuildRegionRequestsSnafu, ConcurrentQueryLimiterClosedSnafu, ConcurrentQueryLimiterTimeoutSnafu, DataFusionSnafu, DecodeLogicalPlanSnafu, @@ -354,7 +355,17 @@ impl RegionServer { None }; - let ctx = request.header.as_ref().map(|h| h.into()); + let ctx = request + .header + .as_ref() + .map(QueryContext::try_from) + .transpose() + .map_err(|err| { + crate::error::InvalidQueryContextSnafu { + reason: err.to_string(), + } + .build() + })?; let query_ctx = Arc::new(ctx.unwrap_or_else(|| QueryContextBuilder::default().build())); let region_id = request.region_id; @@ -977,7 +988,15 @@ impl FlightCraft for RegionServer { let query_ctx = request .header .as_ref() - .map(|h| Arc::new(QueryContext::from(h))) + .map(QueryContext::try_from) + .transpose() + .map_err(|err| { + servers_error::InvalidParameterSnafu { + reason: err.to_string(), + } + .build() + })? + .map(Arc::new) .unwrap_or(QueryContext::arc()); let region_server = self.clone(); @@ -1286,13 +1305,17 @@ impl RegionServerInner { requests: Vec<(RegionId, RegionOpenRequest)>, ignore_nonexistent_region: bool, ) -> Result> { + let request_count = requests.len(); let region_changes = requests .iter() .map(|(region_id, open)| { let attribute = parse_region_attribute(&open.engine, &open.options)?; Ok((*region_id, RegionChange::Register(attribute))) }) - .collect::>>()?; + .collect::>>() + .inspect_err(|err| { + record_region_open_failures(engine.name(), err.status_code(), request_count); + })?; for (®ion_id, region_change) in ®ion_changes { self.set_region_status_not_ready(region_id, &engine, region_change) @@ -1315,6 +1338,7 @@ impl RegionServerInner { .await { error!(e; "Failed to set region to ready: {}", region_id); + record_region_open_failures(engine.name(), e.status_code(), 1); errors.push(BoxedError::new(e)); } else { open_regions.push(region_id) @@ -1328,6 +1352,7 @@ impl RegionServerInner { warn!("Region {} not found, ignore it, source: {:?}", region_id, e); } else { error!(e; "Failed to open region: {}", region_id); + record_region_open_failures(engine.name(), e.status_code(), 1); errors.push(e); } } @@ -1339,6 +1364,7 @@ impl RegionServerInner { self.unset_region_status(region_id, &engine, *region_change); } error!(e; "Failed to open batch regions"); + record_region_open_failures(engine.name(), e.status_code(), request_count); errors.push(BoxedError::new(e)); } } @@ -1374,7 +1400,10 @@ impl RegionServerInner { .read() .unwrap() .get(&engine) - .with_context(|| RegionEngineNotFoundSnafu { name: &engine })? + .with_context(|| RegionEngineNotFoundSnafu { name: &engine }) + .inspect_err(|err| { + record_region_open_failures(&engine, err.status_code(), requests.len()); + })? .clone(); results.push( self.handle_batch_open_requests_inner( @@ -1575,6 +1604,24 @@ impl RegionServerInner { &self, region_id: RegionId, request: RegionRequest, + ) -> Result { + let open_engine = match &request { + RegionRequest::Open(open) => Some(open.engine.clone()), + _ => None, + }; + self.handle_request_inner(region_id, request) + .await + .inspect_err(|err| { + if let Some(engine) = open_engine { + record_region_open_failures(&engine, err.status_code(), 1); + } + }) + } + + async fn handle_request_inner( + &self, + region_id: RegionId, + request: RegionRequest, ) -> Result { let request_type = request.request_type(); let _timer = crate::metrics::HANDLE_REGION_REQUEST_ELAPSED @@ -1921,15 +1968,19 @@ impl RegionServerInner { drop(self.mito_engine.write().unwrap().take()); let engines = self.engines.write().unwrap().drain().collect::>(); + let mut first_error = None; for (engine_name, engine) in engines { - engine + let result = engine .stop() .await - .context(StopRegionEngineSnafu { name: &engine_name })?; - info!("Region engine {engine_name} is stopped"); + .context(StopRegionEngineSnafu { name: &engine_name }); + if result.is_ok() { + info!("Region engine {engine_name} is stopped"); + } + record_shutdown_error(&mut first_error, result); } - Ok(()) + first_error.map_or(Ok(()), Err) } } @@ -1947,6 +1998,17 @@ fn is_metric_engine(engine: &str) -> bool { engine == METRIC_ENGINE_NAME } +/// Records failed open attempts without using arbitrary engine names as metric labels. +fn record_region_open_failures(engine: &str, status_code: StatusCode, count: usize) { + let engine = match engine { + MITO_ENGINE_NAME | METRIC_ENGINE_NAME | FILE_ENGINE_NAME => engine, + _ => "unknown", + }; + crate::metrics::REGION_OPEN_FAILURES_TOTAL + .with_label_values(&[engine, status_code.as_ref()]) + .inc_by(count as u64); +} + fn parse_region_attribute( engine: &str, options: &HashMap, diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs index 9ca771d2186..793ae057076 100644 --- a/src/datanode/src/tests.rs +++ b/src/datanode/src/tests.rs @@ -124,6 +124,8 @@ pub type MockSyncRegionHandler = Box< dyn Fn(RegionId, SyncRegionFromRequest) -> Result + Send + Sync, >; +pub type MockStopHandler = Box Result<(), Error> + Send + Sync>; + pub struct MockRegionEngine { sender: Sender<(RegionId, RegionRequest)>, pub(crate) handle_request_delay: Option, @@ -131,6 +133,7 @@ pub struct MockRegionEngine { pub(crate) handle_set_readonly_gracefully_mock_fn: Option, pub(crate) handle_get_metadata_mock_fn: Option, pub(crate) handle_sync_region_mock_fn: Option, + pub(crate) handle_stop_mock_fn: Option, pub(crate) mock_role: Option>, engine: String, } @@ -147,6 +150,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: None, handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }), @@ -168,6 +172,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: None, handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }), @@ -189,6 +194,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: Some(mock_fn), handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }), @@ -211,6 +217,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: None, handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }; @@ -272,6 +279,9 @@ impl RegionEngine for MockRegionEngine { } async fn stop(&self) -> Result<(), BoxedError> { + if let Some(mock_fn) = &self.handle_stop_mock_fn { + return mock_fn().map_err(BoxedError::new); + } Ok(()) } diff --git a/src/datatypes/Cargo.toml b/src/datatypes/Cargo.toml index 124b1e0ea64..d1dd70acef1 100644 --- a/src/datatypes/Cargo.toml +++ b/src/datatypes/Cargo.toml @@ -37,3 +37,8 @@ serde_json.workspace = true snafu.workspace = true sqlparser.workspace = true sqlparser_derive = "0.1" + +[package.metadata.cargo-shear] +# `sqlparser` is required by `sqlparser_derive`: the generated code references +# `::sqlparser` paths, so it must stay a direct dependency. +ignored = ["sqlparser"] diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs index f2ed05534d6..8ad388c2ff7 100644 --- a/src/datatypes/src/data_type.rs +++ b/src/datatypes/src/data_type.rs @@ -492,7 +492,9 @@ impl TryFrom<&ArrowDataType> for ConcreteDataType { ArrowDataType::Decimal128(precision, scale) => { ConcreteDataType::decimal128_datatype(*precision, *scale) } - ArrowDataType::Struct(fields) => ConcreteDataType::Struct(StructType::from(fields)), + ArrowDataType::Struct(fields) => { + ConcreteDataType::Struct(StructType::try_from_arrow_fields(fields)?) + } ArrowDataType::Float16 | ArrowDataType::Date64 | ArrowDataType::FixedSizeBinary(_) diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 3aa0a4c9792..2e86d7791c5 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -24,9 +24,10 @@ use parquet_variant_compute::VariantType; use serde::{Deserialize, Serialize}; use snafu::{ResultExt, ensure}; -use crate::error::InvalidJson2LayoutSnafu; +use crate::error::{InvalidJson2LayoutSnafu, SerializeSnafu}; pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; +use crate::schema::Metadata; /// Aligns JSON2 field types with their built arrays while preserving field and schema metadata. pub fn align_schema_with_json_array(schema: SchemaRef, columns: &[ArrayRef]) -> SchemaRef { @@ -151,7 +152,7 @@ pub struct JsonMetadata { } impl JsonMetadata { - /// Creates metadata for the JSON2 layout version 2. + /// Creates metadata for the latest JSON2 layout (currently V2). pub fn new(json_settings: JsonSettings) -> Self { Self { json_settings, @@ -300,6 +301,27 @@ impl ExtensionType for Json2ExtensionType { } } +/// Returns JSON2 column metadata with updated settings and the latest layout. +/// +/// Existing SSTs retain their own physical layout metadata. +pub fn json2_metadata_with_updated_settings( + current_metadata: &Metadata, + settings: JsonSettings, +) -> crate::error::Result { + let json_metadata = JsonMetadata::new(settings); + + let mut metadata = current_metadata.clone(); + metadata.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + Json2ExtensionType::NAME.to_string(), + ); + metadata.insert( + EXTENSION_TYPE_METADATA_KEY.to_string(), + serde_json::to_string(&json_metadata).context(SerializeSnafu)?, + ); + Ok(metadata) +} + /// Checks whether this field is either a legacy JSONB or JSON2 extension type. pub fn is_any_json_extension_type>(field: T) -> bool { let name = field.as_ref().extension_type_name(); @@ -442,6 +464,35 @@ mod tests { Ok(()) } + #[test] + fn test_json2_metadata_with_updated_settings_upgrades_layout() -> crate::error::Result<()> { + for (name, json) in [ + (JsonExtensionType::NAME, r#"{"json_settings":{}}"#), + (Json2ExtensionType::NAME, r#"{"json_settings":{}}"#), + ( + Json2ExtensionType::NAME, + r#"{"json_settings":{},"layout_version":2}"#, + ), + ] { + let metadata = HashMap::from([ + (EXTENSION_TYPE_NAME_KEY.to_string(), name.to_string()), + (EXTENSION_TYPE_METADATA_KEY.to_string(), json.to_string()), + ("other".to_string(), "kept".to_string()), + ]); + let settings = JsonSettings::try_new(vec![], Some(10))?; + let updated = json2_metadata_with_updated_settings(&metadata, settings.clone())?; + assert_eq!(Some("kept"), updated.get("other").map(String::as_str)); + assert_eq!( + Some(Json2ExtensionType::NAME), + updated.get(EXTENSION_TYPE_NAME_KEY).map(String::as_str) + ); + let json_metadata: JsonMetadata = + serde_json::from_str(updated.get(EXTENSION_TYPE_METADATA_KEY).unwrap()).unwrap(); + assert_eq!(JsonMetadata::new(settings), json_metadata); + } + Ok(()) + } + #[test] fn test_parse_json2_physical_layout() -> crate::error::Result<()> { let legacy = Field::new("data", DataType::Struct(Fields::empty()), true) diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index a5b4940e048..e3c408698eb 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -23,6 +23,7 @@ pub mod value; use std::collections::BTreeMap; use std::collections::btree_map::Entry; +use std::sync::Arc; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; @@ -31,6 +32,7 @@ use snafu::ResultExt; use crate::data_type::ConcreteDataType; use crate::error::{self, InvalidJson2SettingsSnafu, Result, UnsupportedJsonTypeSnafu}; use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb}; +use crate::prelude::DataType as _; use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::value::{ListValue, StructValue, Value}; @@ -93,6 +95,15 @@ pub struct JsonTypeHint { pub inverted_index: bool, } +/// Specifies how encoding handles a value that does not match its JSON2 type hint. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TypeHintMismatchPolicy { + /// Returns an error for a type hint mismatch. + Reject, + /// Coerces to the hint type, encoding unconvertible values as null. + CoerceOrNull, +} + /// Context for JSON encoding/decoding that tracks the current key path. #[derive(Clone, Debug)] pub struct JsonContext<'a> { @@ -135,6 +146,30 @@ impl JsonSettings { (self.type_hints, self.max_auto_expanded_paths) } + /// Returns whether the settings are equivalent, ignoring type hint order. + pub fn equivalent(&self, other: &Self) -> bool { + let Self { + type_hints, + max_auto_expanded_paths, + } = self; + let Self { + type_hints: other_hints, + max_auto_expanded_paths: other_max_auto_expanded_paths, + } = other; + + if max_auto_expanded_paths != other_max_auto_expanded_paths + || type_hints.len() != other_hints.len() + { + return false; + } + + let mut hints = type_hints.iter().collect::>(); + let mut other_hints = other_hints.iter().collect::>(); + hints.sort_unstable_by(|a, b| a.path.cmp(&b.path)); + other_hints.sort_unstable_by(|a, b| a.path.cmp(&b.path)); + hints == other_hints + } + /// Decode an encoded StructValue back into a serde_json::Value. pub fn decode(&self, value: Value) -> Result { let mut context = JsonContext { @@ -146,6 +181,15 @@ impl JsonSettings { /// Encode a serde_json::Value into a Value::Json using current settings. pub fn encode(&self, json: Json) -> Result { + self.encode_with_type_hint_mismatch_policy(json, TypeHintMismatchPolicy::Reject) + } + + /// Encodes a serde_json::Value using the given type hint mismatch policy. + pub fn encode_with_type_hint_mismatch_policy( + &self, + json: Json, + policy: TypeHintMismatchPolicy, + ) -> Result { if let Json::Object(object) = &json && object.contains_key(JSON2_REMAINDER_FIELD_NAME) { @@ -160,7 +204,7 @@ impl JsonSettings { path: Vec::new(), settings: self, }; - encode_json_with_context(json, &mut context).map(|v| Value::Json(Box::new(v))) + encode_json_with_context(json, &mut context, policy).map(|v| Value::Json(Box::new(v))) } } @@ -193,19 +237,10 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { } let data_type = match &hint.data_type { ConcreteDataType::Boolean(_) - | ConcreteDataType::UInt8(_) - | ConcreteDataType::UInt16(_) - | ConcreteDataType::UInt32(_) | ConcreteDataType::UInt64(_) - | ConcreteDataType::Int8(_) - | ConcreteDataType::Int16(_) - | ConcreteDataType::Int32(_) | ConcreteDataType::Int64(_) - | ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_) - | ConcreteDataType::String(_) - | ConcreteDataType::List(_) - | ConcreteDataType::Struct(_) => (&hint.data_type).into(), + | ConcreteDataType::String(_) => (&hint.data_type).into(), data_type => { return InvalidJson2SettingsSnafu { reason: format!("unsupported JSON2 type hint data type: {data_type}"), @@ -272,36 +307,41 @@ fn with_key_context( } /// Main encoding function with key path tracking -fn encode_json_with_context(json: Json, context: &mut JsonContext) -> Result { +fn encode_json_with_context( + json: Json, + context: &mut JsonContext, + policy: TypeHintMismatchPolicy, +) -> Result { if context.path.is_empty() && !matches!(json, Json::Object(_)) { return UnsupportedJsonTypeSnafu.fail(); } match json { - Json::Object(json_object) => encode_json_object_with_context(json_object, context), - Json::Array(json_array) => encode_json_array_with_context(json_array, context), - _ => encode_json_value_with_context(json, context), + Json::Object(json_object) => encode_json_object_with_context(json_object, context, policy), + Json::Array(json_array) => encode_json_array_with_context(json_array, context, policy), + _ => encode_json_value_with_context(json, context, policy), } } fn encode_json_object_with_context<'a>( json_object: Map, context: &mut JsonContext<'a>, + policy: TypeHintMismatchPolicy, ) -> Result { let mut object = BTreeMap::new(); for (key, value) in json_object { let value = with_key_context(context, &key, |context| { if let Some(hint) = context.type_hint() { - encode_json_value_with_hint(value, hint, context) + encode_json_value_with_hint(value, hint, context, policy) } else { - encode_json_value_with_context(value, context) + encode_json_value_with_context(value, context, policy) } })?; object.insert(key, value.into_variant()); } - fill_missing_type_hints(&mut object, context)?; + fill_missing_type_hints(&mut object, context, policy)?; Ok(JsonValue::new(JsonVariant::Object(object))) } @@ -309,13 +349,14 @@ fn encode_json_object_with_context<'a>( fn fill_missing_type_hints( object: &mut BTreeMap, context: &mut JsonContext, + policy: TypeHintMismatchPolicy, ) -> Result<()> { for hint in &context.settings.type_hints { if hint.path.len() > context.path.len() && hint.path.starts_with(&context.path) { let depth = context.path.len(); let key = &hint.path[depth]; with_key_context(context, key, |context| { - insert_missing_type_hint(object, context, hint, depth) + insert_missing_type_hint(object, context, hint, depth, policy) })?; } } @@ -327,6 +368,7 @@ fn insert_missing_type_hint( field_context: &mut JsonContext, hint: &JsonTypeHint, depth: usize, + policy: TypeHintMismatchPolicy, ) -> Result<()> { let key = &hint.path[depth]; let is_leaf = depth + 1 == hint.path.len(); @@ -341,7 +383,11 @@ fn insert_missing_type_hint( match object.entry(key.clone()) { Entry::Occupied(mut entry) => match entry.get_mut() { JsonVariant::Object(child) => { - insert_missing_type_hint(child, field_context, hint, depth + 1) + insert_missing_type_hint(child, field_context, hint, depth + 1, policy) + } + _ if policy == TypeHintMismatchPolicy::CoerceOrNull => { + entry.insert(JsonValue::null().into_variant()); + Ok(()) } _ => error::InvalidJsonSnafu { value: format!( @@ -354,7 +400,7 @@ fn insert_missing_type_hint( }, Entry::Vacant(entry) => { let mut child = BTreeMap::new(); - insert_missing_type_hint(&mut child, field_context, hint, depth + 1)?; + insert_missing_type_hint(&mut child, field_context, hint, depth + 1, policy)?; entry.insert(JsonVariant::Object(child)); Ok(()) } @@ -365,65 +411,152 @@ fn encode_json_value_with_hint( json: Json, hint: &JsonTypeHint, context: &mut JsonContext, + policy: TypeHintMismatchPolicy, ) -> Result { if json.is_null() { return Ok(JsonValue::null()); } - let invalid_type = || { - error::InvalidJsonSnafu { + let json = match (&hint.data_type, json) { + (ConcreteDataType::String(_), Json::String(value)) => { + return Ok(value.into()); + } + ( + ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_), + Json::Number(value), + ) => { + if let Some(value) = value.as_i64() { + return Ok(value.into()); + } + Json::Number(value) + } + ( + ConcreteDataType::UInt8(_) + | ConcreteDataType::UInt16(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::UInt64(_), + Json::Number(value), + ) => { + if let Some(value) = value.as_u64() { + return Ok(value.into()); + } + Json::Number(value) + } + (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), Json::Number(value)) => { + if let Some(value) = value.as_f64() { + return Ok(value.into()); + } + Json::Number(value) + } + (ConcreteDataType::Boolean(_), Json::Bool(value)) => { + return Ok(value.into()); + } + (_, json) => json, + }; + + match policy { + TypeHintMismatchPolicy::Reject => error::InvalidJsonSnafu { value: format!( "JSON value at {} does not match JSON2 type hint {}", context.path.join("."), hint.data_type ), } - .fail() - }; + .fail(), + TypeHintMismatchPolicy::CoerceOrNull => { + let value = coerce_json_value_to_type(json, &hint.data_type); + let value = Json::try_from(value).map_err(|error| { + error::InvalidJsonSnafu { + value: error.to_string(), + } + .build() + })?; + Ok(JsonValue::new(value.into())) + } + } +} - match (&hint.data_type, json) { - (ConcreteDataType::String(_), Json::String(v)) => Ok(v.into()), - ( - ConcreteDataType::Int8(_) - | ConcreteDataType::Int16(_) - | ConcreteDataType::Int32(_) - | ConcreteDataType::Int64(_), - Json::Number(v), - ) => match v.as_i64() { - Some(v) => Ok(v.into()), - None => invalid_type(), - }, - ( - ConcreteDataType::UInt8(_) - | ConcreteDataType::UInt16(_) - | ConcreteDataType::UInt32(_) - | ConcreteDataType::UInt64(_), - Json::Number(v), - ) => match v.as_u64() { - Some(v) => Ok(v.into()), - None => invalid_type(), - }, - (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), Json::Number(v)) => { - match v.as_f64() { - Some(v) => Ok(v.into()), - None => invalid_type(), +/// Coerces a JSON value using the same semantics as JSON2 query projection. +/// Values that cannot be converted become null. +pub(crate) fn coerce_json_value_to_type(value: Json, to_type: &ConcreteDataType) -> Value { + if value.is_null() { + return Value::Null; + } + + if to_type.is_string() { + let value = match value { + Json::String(value) => value, + value => value.to_string(), + }; + return Value::String(value.into()); + } + + if matches!(to_type, ConcreteDataType::Binary(_)) { + return Value::Binary(encode_serde_json_as_jsonb(value).into()); + } + + if let Some(struct_type) = to_type.as_struct() { + let Json::Object(mut object) = value else { + return Value::Null; + }; + let values = struct_type + .fields() + .iter() + .map(|field| { + object + .remove(field.name()) + .map(|value| coerce_json_value_to_type(value, field.data_type())) + .unwrap_or(Value::Null) + }) + .collect::>(); + return Value::Struct(StructValue::new(values, struct_type.clone())); + } + + if let Some(list_type) = to_type.as_list() { + let Json::Array(values) = value else { + return Value::Null; + }; + let item_type = list_type.item_type().clone(); + let values = values + .into_iter() + .map(|value| coerce_json_value_to_type(value, &item_type)) + .collect::>(); + return Value::List(ListValue::new(values, Arc::new(item_type))); + } + + let value = match value { + Json::Bool(value) => Value::Boolean(value), + Json::Number(value) => { + if let Some(value) = value.as_i64() { + Value::Int64(value) + } else if let Some(value) = value.as_u64() { + Value::UInt64(value) + } else if let Some(value) = value.as_f64() { + Value::Float64(value.into()) + } else { + Value::Null } } - (ConcreteDataType::Boolean(_), Json::Bool(v)) => Ok(v.into()), - _ => invalid_type(), - } + Json::String(value) => Value::String(value.into()), + Json::Array(_) | Json::Object(_) | Json::Null => Value::Null, + }; + to_type.try_cast(value).unwrap_or(Value::Null) } fn encode_json_array_with_context<'a>( json_array: Vec, context: &mut JsonContext<'a>, + policy: TypeHintMismatchPolicy, ) -> Result { let json_array_len = json_array.len(); let mut items = Vec::with_capacity(json_array_len); for (index, value) in json_array.into_iter().enumerate() { let item_value = with_key_context(context, &index.to_string(), |context| { - encode_json_value_with_context(value, context) + encode_json_value_with_context(value, context, policy) })?; items.push(item_value); } @@ -458,7 +591,11 @@ fn encode_json_array_with_context<'a>( } /// Helper function to encode a JSON value to a Value and determine its ConcreteDataType with context -fn encode_json_value_with_context(json: Json, context: &mut JsonContext) -> Result { +fn encode_json_value_with_context( + json: Json, + context: &mut JsonContext, + policy: TypeHintMismatchPolicy, +) -> Result { if context.path.len() >= JSON2_MAX_STRUCTURED_DEPTH && matches!(&json, Json::Object(_) | Json::Array(_)) { @@ -487,8 +624,8 @@ fn encode_json_value_with_context(json: Json, context: &mut JsonContext) -> Resu } } Json::String(s) => Ok(s.into()), - Json::Array(arr) => encode_json_array_with_context(arr, context), - Json::Object(obj) => encode_json_object_with_context(obj, context), + Json::Array(arr) => encode_json_array_with_context(arr, context, policy), + Json::Object(obj) => encode_json_object_with_context(obj, context, policy), } } @@ -597,6 +734,50 @@ mod tests { &struct_value.items()[index] } + #[test] + fn test_json_settings_equivalent() -> Result<()> { + let hints = vec![ + JsonTypeHint { + path: vec!["user".to_string(), "name".to_string()], + data_type: ConcreteDataType::string_datatype(), + inverted_index: false, + }, + JsonTypeHint { + path: vec!["count".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }, + ]; + let settings = JsonSettings::try_new(hints.clone(), Some(10))?; + let mut reversed = hints.clone(); + reversed.reverse(); + let reordered = JsonSettings::try_new(reversed, Some(10))?; + assert_ne!(settings, reordered); + assert!(settings.equivalent(&reordered)); + assert!(reordered.equivalent(&settings)); + assert!(settings.equivalent(&settings)); + assert_eq!(hints, settings.type_hints()); + assert!(JsonSettings::default().equivalent(&JsonSettings::default())); + + for limit in [None, Some(0), Some(11)] { + assert!(!settings.equivalent(&JsonSettings::try_new(hints.clone(), limit)?)); + } + assert!(!settings.equivalent(&JsonSettings::try_new(vec![hints[0].clone()], Some(10))?)); + + let mut changed = hints.clone(); + changed[0].path = vec!["user".to_string(), "id".to_string()]; + assert!(!settings.equivalent(&JsonSettings::try_new(changed, Some(10))?)); + + let mut changed = hints.clone(); + changed[0].data_type = ConcreteDataType::int64_datatype(); + assert!(!settings.equivalent(&JsonSettings::try_new(changed, Some(10))?)); + + let mut changed = hints; + changed[0].inverted_index = true; + assert!(!settings.equivalent(&JsonSettings::try_new(changed, Some(10))?)); + Ok(()) + } + #[test] fn test_json_settings_deserializes_legacy_type_hint_constraints() -> std::result::Result<(), Box> { @@ -755,6 +936,104 @@ mod tests { } } + #[test] + fn test_coerce_or_null_on_type_hint_mismatch() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }], + None, + )?; + + let result = settings + .encode_with_type_hint_mismatch_policy( + json!({"kind": "invalid", "message": "kept"}), + TypeHintMismatchPolicy::CoerceOrNull, + )? + .into_json_inner() + .unwrap(); + let Value::Struct(root) = result else { + panic!("Expected Struct value"); + }; + assert_eq!(struct_field_value(&root, "kind"), &Value::Null); + assert_eq!( + struct_field_value(&root, "message"), + &Value::String("kept".into()) + ); + Ok(()) + } + + #[test] + fn test_coerce_or_null_converts_type_hint_values() -> Result<()> { + let settings = JsonSettings::try_new( + vec![ + JsonTypeHint { + path: vec!["to_string".to_string()], + data_type: ConcreteDataType::string_datatype(), + inverted_index: false, + }, + JsonTypeHint { + path: vec!["to_int".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }, + ], + None, + )?; + + let result = settings + .encode_with_type_hint_mismatch_policy( + json!({"to_string": 123, "to_int": "456", "invalid": "value"}), + TypeHintMismatchPolicy::CoerceOrNull, + )? + .into_json_inner() + .unwrap(); + let Value::Struct(root) = result else { + panic!("Expected Struct value"); + }; + assert_eq!( + struct_field_value(&root, "to_string"), + &Value::String("123".into()) + ); + assert_eq!(struct_field_value(&root, "to_int"), &Value::Int64(456)); + assert_eq!( + struct_field_value(&root, "invalid"), + &Value::String("value".into()) + ); + Ok(()) + } + + #[test] + fn test_coerce_or_null_on_nested_type_hint_structure_mismatch() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["user".to_string(), "age".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }], + None, + )?; + + let result = settings + .encode_with_type_hint_mismatch_policy( + json!({"user": "invalid", "message": "kept"}), + TypeHintMismatchPolicy::CoerceOrNull, + )? + .into_json_inner() + .unwrap(); + let Value::Struct(root) = result else { + panic!("Expected Struct value"); + }; + assert_eq!(struct_field_value(&root, "user"), &Value::Null); + assert_eq!( + struct_field_value(&root, "message"), + &Value::String("kept".into()) + ); + Ok(()) + } + #[test] fn test_encode_rejects_reserved_remainder_field() -> Result<()> { let settings = JsonSettings::default(); diff --git a/src/datatypes/src/schema.rs b/src/datatypes/src/schema.rs index d53b42e3e7e..0f1a832ebd2 100644 --- a/src/datatypes/src/schema.rs +++ b/src/datatypes/src/schema.rs @@ -32,13 +32,9 @@ pub use crate::schema::column_schema::{ COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, - COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, - COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, COMMENT_KEY, - ColumnExtType, ColumnSchema, FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, - INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, - TIME_INDEX_KEY, VECTOR_INDEX_KEY, VectorDistanceMetric, VectorIndexEngineType, - VectorIndexOptions, + COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, ColumnExtType, ColumnSchema, FULLTEXT_KEY, + FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, + SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, TIME_INDEX_KEY, }; pub use crate::schema::constraint::ColumnDefaultConstraint; diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index a71605afdb2..b3ce9cafedb 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -48,8 +48,6 @@ pub const FULLTEXT_KEY: &str = "greptime:fulltext"; pub const INVERTED_INDEX_KEY: &str = "greptime:inverted_index"; /// Key used to store skip options in arrow field's metadata. pub const SKIPPING_INDEX_KEY: &str = "greptime:skipping_index"; -/// Key used to store vector index options in arrow field's metadata. -pub const VECTOR_INDEX_KEY: &str = "greptime:vector_index"; /// Keys used in fulltext options pub const COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE: &str = "enable"; @@ -64,13 +62,6 @@ pub const COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY: &str = "granularity"; pub const COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE: &str = "false_positive_rate"; pub const COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE: &str = "type"; -/// Keys used in VECTOR index options -pub const COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE: &str = "engine"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_METRIC: &str = "metric"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY: &str = "connectivity"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD: &str = "expansion_add"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH: &str = "expansion_search"; - pub const DEFAULT_GRANULARITY: u32 = 10240; pub const DEFAULT_FALSE_POSITIVE_RATE: f64 = 0.01; @@ -255,53 +246,6 @@ impl ColumnSchema { self.metadata.contains_key(INVERTED_INDEX_KEY) } - /// Checks if this column has a vector index. - pub fn is_vector_indexed(&self) -> bool { - match self.vector_index_options() { - Ok(opts) => opts.is_some(), - Err(e) => { - common_telemetry::warn!( - "Failed to deserialize vector_index_options for column '{}': {}", - self.name, - e - ); - false - } - } - } - - /// Gets the vector index options. - pub fn vector_index_options(&self) -> Result> { - match self.metadata.get(VECTOR_INDEX_KEY) { - None => Ok(None), - Some(json) => { - let options = - serde_json::from_str(json).context(error::DeserializeSnafu { json })?; - Ok(Some(options)) - } - } - } - - /// Sets the vector index options. - pub fn set_vector_index_options(&mut self, options: &VectorIndexOptions) -> Result<()> { - self.metadata.insert( - VECTOR_INDEX_KEY.to_string(), - serde_json::to_string(options).context(error::SerializeSnafu)?, - ); - Ok(()) - } - - /// Removes the vector index options. - pub fn unset_vector_index_options(&mut self) { - self.metadata.remove(VECTOR_INDEX_KEY); - } - - /// Sets vector index options and returns self for chaining. - pub fn with_vector_index_options(mut self, options: &VectorIndexOptions) -> Result { - self.set_vector_index_options(options)?; - Ok(self) - } - /// Set default constraint. /// /// If a default constraint exists for the column, this method will @@ -1067,181 +1011,6 @@ impl TryFrom> for SkippingIndexOptions { } } -/// Distance metric for vector similarity search. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, Visit, VisitMut)] -#[serde(rename_all = "lowercase")] -pub enum VectorDistanceMetric { - /// Squared Euclidean distance (L2^2). - #[default] - L2sq, - /// Cosine distance (1 - cosine similarity). - Cosine, - /// Inner product (negative, for maximum inner product search). - #[serde(alias = "ip")] - InnerProduct, -} - -impl fmt::Display for VectorDistanceMetric { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - VectorDistanceMetric::L2sq => write!(f, "l2sq"), - VectorDistanceMetric::Cosine => write!(f, "cosine"), - VectorDistanceMetric::InnerProduct => write!(f, "ip"), - } - } -} - -impl std::str::FromStr for VectorDistanceMetric { - type Err = String; - - fn from_str(s: &str) -> std::result::Result { - match s.to_lowercase().as_str() { - "l2sq" | "l2" | "euclidean" => Ok(VectorDistanceMetric::L2sq), - "cosine" | "cos" => Ok(VectorDistanceMetric::Cosine), - "inner_product" | "ip" | "dot" => Ok(VectorDistanceMetric::InnerProduct), - _ => Err(format!( - "Unknown distance metric: {}. Expected: l2sq, cosine, or ip", - s - )), - } - } -} - -impl VectorDistanceMetric { - /// Returns the metric as u8 for blob serialization. - pub fn as_u8(&self) -> u8 { - match self { - Self::L2sq => 0, - Self::Cosine => 1, - Self::InnerProduct => 2, - } - } - - /// Parses metric from u8 (used when reading blob). - pub fn try_from_u8(v: u8) -> Option { - match v { - 0 => Some(Self::L2sq), - 1 => Some(Self::Cosine), - 2 => Some(Self::InnerProduct), - _ => None, - } - } -} - -/// Default HNSW connectivity parameter. -const DEFAULT_VECTOR_INDEX_CONNECTIVITY: u32 = 16; -/// Default expansion factor during index construction. -const DEFAULT_VECTOR_INDEX_EXPANSION_ADD: u32 = 128; -/// Default expansion factor during search. -const DEFAULT_VECTOR_INDEX_EXPANSION_SEARCH: u32 = 64; - -fn default_vector_index_connectivity() -> u32 { - DEFAULT_VECTOR_INDEX_CONNECTIVITY -} - -fn default_vector_index_expansion_add() -> u32 { - DEFAULT_VECTOR_INDEX_EXPANSION_ADD -} - -fn default_vector_index_expansion_search() -> u32 { - DEFAULT_VECTOR_INDEX_EXPANSION_SEARCH -} - -/// Supported vector index engine types. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, Visit, VisitMut)] -#[serde(rename_all = "lowercase")] -pub enum VectorIndexEngineType { - /// USearch HNSW implementation. - #[default] - Usearch, - // Future: Vsag, -} - -impl VectorIndexEngineType { - /// Returns the engine type as u8 for blob serialization. - pub fn as_u8(&self) -> u8 { - match self { - Self::Usearch => 0, - } - } - - /// Parses engine type from u8 (used when reading blob). - pub fn try_from_u8(v: u8) -> Option { - match v { - 0 => Some(Self::Usearch), - _ => None, - } - } -} - -impl fmt::Display for VectorIndexEngineType { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Usearch => write!(f, "usearch"), - } - } -} - -impl std::str::FromStr for VectorIndexEngineType { - type Err = String; - - fn from_str(s: &str) -> std::result::Result { - match s.to_lowercase().as_str() { - "usearch" => Ok(Self::Usearch), - _ => Err(format!( - "Unknown vector index engine: {}. Expected: usearch", - s - )), - } - } -} - -/// Options for vector index (HNSW). -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)] -#[serde(rename_all = "kebab-case")] -pub struct VectorIndexOptions { - /// Vector index engine type (default: usearch). - #[serde(default)] - pub engine: VectorIndexEngineType, - /// Distance metric for similarity search. - #[serde(default)] - pub metric: VectorDistanceMetric, - /// HNSW connectivity parameter (M in the paper). - /// Higher values improve recall but increase memory usage. - #[serde(default = "default_vector_index_connectivity")] - pub connectivity: u32, - /// Expansion factor during index construction (ef_construction). - /// Higher values improve index quality but slow down construction. - #[serde(default = "default_vector_index_expansion_add")] - pub expansion_add: u32, - /// Expansion factor during search (ef_search). - /// Higher values improve recall but slow down search. - #[serde(default = "default_vector_index_expansion_search")] - pub expansion_search: u32, -} - -impl Default for VectorIndexOptions { - fn default() -> Self { - Self { - engine: VectorIndexEngineType::default(), - metric: VectorDistanceMetric::default(), - connectivity: DEFAULT_VECTOR_INDEX_CONNECTIVITY, - expansion_add: DEFAULT_VECTOR_INDEX_EXPANSION_ADD, - expansion_search: DEFAULT_VECTOR_INDEX_EXPANSION_SEARCH, - } - } -} - -impl fmt::Display for VectorIndexOptions { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "engine={}, metric={}, connectivity={}, expansion_add={}, expansion_search={}", - self.engine, self.metric, self.connectivity, self.expansion_add, self.expansion_search - ) - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/datatypes/src/types/struct_type.rs b/src/datatypes/src/types/struct_type.rs index ffd1d1c0b25..fdbfd4886b9 100644 --- a/src/datatypes/src/types/struct_type.rs +++ b/src/datatypes/src/types/struct_type.rs @@ -19,6 +19,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Field}; use arrow_schema::Fields; use serde::{Deserialize, Serialize}; +use crate::error::Result; use crate::prelude::{ConcreteDataType, DataType, LogicalTypeId}; use crate::value::Value; use crate::vectors::StructVectorBuilder; @@ -100,6 +101,24 @@ impl StructType { } } + /// Like the `From<&Fields>` impl, but returns an error for arrow field + /// types greptimedb cannot represent instead of panicking. + pub fn try_from_arrow_fields(fields: &Fields) -> Result { + let struct_fields = fields + .iter() + .map(|field| { + Ok(StructField::new( + field.name().clone(), + ConcreteDataType::try_from(field.data_type())?, + field.is_nullable(), + )) + }) + .collect::>>()?; + Ok(StructType { + fields: Arc::new(struct_fields), + }) + } + pub fn fields(&self) -> Arc> { self.fields.clone() } diff --git a/src/datatypes/src/types/vector_type.rs b/src/datatypes/src/types/vector_type.rs index d01949c4737..ab756be90eb 100644 --- a/src/datatypes/src/types/vector_type.rs +++ b/src/datatypes/src/types/vector_type.rs @@ -86,8 +86,10 @@ pub fn vector_type_value_to_string(val: &[u8], dim: u32) -> Result { } let elements = val - .chunks_exact(std::mem::size_of::()) - .map(|e| f32::from_le_bytes(e.try_into().unwrap())); + .as_chunks::<{ std::mem::size_of::() }>() + .0 + .iter() + .map(|e| f32::from_le_bytes(*e)); let mut s = String::from("["); for (i, e) in elements.enumerate() { diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 56a551484f2..da1df882aba 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -1152,12 +1152,13 @@ impl TryFrom for Value { ScalarValue::UInt16(u) => Value::from(u), ScalarValue::UInt32(u) => Value::from(u), ScalarValue::UInt64(u) => Value::from(u), - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) => { + ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s) => { Value::from(s.map(StringBytes::from)) } ScalarValue::Binary(b) | ScalarValue::LargeBinary(b) - | ScalarValue::FixedSizeBinary(_, b) => Value::from(b.map(Bytes::from)), + | ScalarValue::FixedSizeBinary(_, b) + | ScalarValue::BinaryView(b) => Value::from(b.map(Bytes::from)), ScalarValue::List(array) => { // this is for item type let datatype = ConcreteDataType::try_from(&array.value_type())?; @@ -1221,7 +1222,15 @@ impl TryFrom for Value { .map(|v| Value::Decimal128(Decimal128::new(v, p, s))) .unwrap_or(Value::Null), ScalarValue::Struct(struct_array) => { - let struct_type = StructType::from(struct_array.fields()); + // A struct scalar carries a single element; a null (or empty) + // element means the struct itself is null, for example a null + // struct inside a list. + if struct_array.is_empty() || struct_array.is_null(0) { + return Ok(Value::Null); + } + // Build the struct type fallibly: an unrepresentable arrow + // field type must surface as an error, not a panic. + let struct_type = StructType::try_from_arrow_fields(struct_array.fields())?; let items = struct_array .columns() .iter() @@ -1244,8 +1253,6 @@ impl TryFrom for Value { | ScalarValue::LargeListView(_) | ScalarValue::Union(_, _, _) | ScalarValue::Float16(_) - | ScalarValue::Utf8View(_) - | ScalarValue::BinaryView(_) | ScalarValue::Map(_) | ScalarValue::Date64(_) | ScalarValue::RunEndEncoded(_, _, _) => { @@ -1259,6 +1266,21 @@ impl TryFrom for Value { } } +/// Converts the value at `index` of an arrow array into a [`Value`]. +/// +/// Unlike the unchecked conversion in [`crate::vectors::StructVector`]'s `get`, +/// this propagates conversion errors (for example arrow types greptimedb +/// cannot represent, such as `Decimal256`) instead of panicking, so callers +/// can surface a proper error to the user. +pub fn try_value_from_array(array: &dyn Array, index: usize) -> Result { + if array.is_null(index) { + return Ok(Value::Null); + } + let scalar = + ScalarValue::try_from_array(array, index).context(ConvertArrowArrayToScalarsSnafu)?; + Value::try_from(scalar) +} + impl From> for Value { fn from(value: ValueRef<'_>) -> Self { match value { @@ -2157,6 +2179,73 @@ pub(crate) mod tests { Value::Struct(struct_value), scalar_struct_value.try_into().unwrap() ); + + // view-typed strings and binaries convert like their non-view forms + assert_eq!( + Value::String("abc".into()), + ScalarValue::Utf8View(Some("abc".into())) + .try_into() + .unwrap() + ); + assert_eq!(Value::Null, ScalarValue::Utf8View(None).try_into().unwrap()); + assert_eq!( + Value::Binary(Bytes::from(vec![1, 2])), + ScalarValue::BinaryView(Some(vec![1, 2])) + .try_into() + .unwrap() + ); + assert_eq!( + Value::Null, + ScalarValue::BinaryView(None).try_into().unwrap() + ); + + // a null struct scalar is a null value, not a struct of null fields + let null_struct = ScalarStructBuilder::new_null(build_struct_type().as_arrow_fields()); + assert_eq!(Value::Null, null_struct.try_into().unwrap()); + } + + #[test] + fn test_try_value_from_array() { + use arrow_array::Int32Array; + use datafusion_common::arrow::datatypes::i256; + + let array = Int32Array::from(vec![Some(1), None]); + assert_eq!(Value::Int32(1), try_value_from_array(&array, 0).unwrap()); + assert_eq!(Value::Null, try_value_from_array(&array, 1).unwrap()); + + // a struct with representable fields converts + let supported = ScalarStructBuilder::new() + .with_name_and_scalar("i", ScalarValue::Int32(Some(7))) + .build() + .unwrap(); + let ScalarValue::Struct(array) = supported else { + unreachable!(); + }; + assert_eq!( + Value::Struct(StructValue::new( + vec![Value::Int32(7)], + StructType::new(std::sync::Arc::new(vec![StructField::new( + "i", + ConcreteDataType::int32_datatype(), + false, + )])) + )), + try_value_from_array(array.as_ref(), 0).unwrap() + ); + + // an unsupported arrow field type errors instead of panicking + let unsupported = ScalarStructBuilder::new() + .with_name_and_scalar( + "d", + ScalarValue::Decimal256(Some(i256::from_i128(1)), 38, 10), + ) + .build() + .unwrap(); + let ScalarValue::Struct(array) = unsupported else { + unreachable!(); + }; + let error = try_value_from_array(array.as_ref(), 0).unwrap_err(); + assert!(error.to_string().contains("Unsupported arrow data type")); } #[test] diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 08a2fd86555..5b6265c7b37 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -26,15 +26,13 @@ use serde_json::Value; use snafu::{OptionExt, ResultExt}; use crate::arrow_array::{binary_array_value, string_array_value}; -use crate::data_type::ConcreteDataType; +use crate::data_type::{ConcreteDataType, DataType as _}; use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, }; use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field}; -use crate::json::JsonSettings; -use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb}; -use crate::prelude::{DataType as _, Value as GreptimeValue}; -use crate::value::{ListValue, StructValue}; +use crate::json::value::decode_json_variant; +use crate::json::{JsonSettings, TypeHintMismatchPolicy, coerce_json_value_to_type}; use crate::vectors::MutableVector; use crate::vectors::json::builder::{JsonVectorBuilder, json2_physical_data_type}; use crate::vectors::json::variant::variant_to_json_values; @@ -123,6 +121,23 @@ impl JsonArray<'_> { field: &Field, logical_settings: &JsonSettings, target_layout: &JsonSettings, + ) -> Result { + self.rewrite_to_v2_with_type_hint_mismatch_policy( + field, + logical_settings, + target_layout, + TypeHintMismatchPolicy::Reject, + ) + } + + /// Rewrites a JSON2 array to the specified v2 physical layout using the + /// given type hint mismatch policy. + pub fn rewrite_to_v2_with_type_hint_mismatch_policy( + &self, + field: &Field, + logical_settings: &JsonSettings, + target_layout: &JsonSettings, + policy: TypeHintMismatchPolicy, ) -> Result { let is_v2 = json2_remainder_field(field)?.is_some(); if is_v2 && self.inner.data_type() == &json2_physical_data_type(target_layout) { @@ -141,7 +156,8 @@ impl JsonArray<'_> { if value.is_null() { builder.push_null(); } else { - let value = logical_settings.encode(value)?; + let value = + logical_settings.encode_with_type_hint_mismatch_policy(value, policy)?; builder.try_push_value_ref(&value.as_value_ref())?; } } @@ -319,87 +335,12 @@ fn project_json_values(values: Vec, to_type: &DataType) -> Result Result { - if value.is_null() { - return Ok(GreptimeValue::Null); - } - - if to_type.is_string() { - let value = match value { - Value::String(value) => value, - value => value.to_string(), - }; - return Ok(GreptimeValue::String(value.into())); - } - - if matches!(to_type, ConcreteDataType::Binary(_)) { - return Ok(GreptimeValue::Binary( - encode_serde_json_as_jsonb(value).into(), - )); - } - - if let Some(struct_type) = to_type.as_struct() { - let Value::Object(mut object) = value else { - return Ok(GreptimeValue::Null); - }; - let values = struct_type - .fields() - .iter() - .map(|field| { - object - .remove(field.name()) - .map(|value| project_json_value_to_type(value, field.data_type())) - .transpose() - .map(|value| value.unwrap_or(GreptimeValue::Null)) - }) - .collect::>>()?; - return Ok(GreptimeValue::Struct(StructValue::new( - values, - struct_type.clone(), - ))); - } - - if let Some(list_type) = to_type.as_list() { - let Value::Array(values) = value else { - return Ok(GreptimeValue::Null); - }; - let item_type = list_type.item_type().clone(); - let values = values - .into_iter() - .map(|value| project_json_value_to_type(value, &item_type)) - .collect::>>()?; - return Ok(GreptimeValue::List(ListValue::new( - values, - Arc::new(item_type), - ))); - } - - let value = match value { - Value::Bool(value) => GreptimeValue::Boolean(value), - Value::Number(value) => { - if let Some(value) = value.as_i64() { - GreptimeValue::Int64(value) - } else if let Some(value) = value.as_u64() { - GreptimeValue::UInt64(value) - } else if let Some(value) = value.as_f64() { - GreptimeValue::Float64(value.into()) - } else { - GreptimeValue::Null - } - } - Value::String(value) => GreptimeValue::String(value.into()), - Value::Array(_) | Value::Object(_) => GreptimeValue::Null, - Value::Null => GreptimeValue::Null, - }; - Ok(to_type.try_cast(value).unwrap_or(GreptimeValue::Null)) -} - impl<'a> From<&'a ArrayRef> for JsonArray<'a> { fn from(inner: &'a ArrayRef) -> Self { Self { inner } diff --git a/src/file-engine/Cargo.toml b/src/file-engine/Cargo.toml index ec904addf1e..dda5fb4221c 100644 --- a/src/file-engine/Cargo.toml +++ b/src/file-engine/Cargo.toml @@ -18,7 +18,6 @@ common-catalog.workspace = true common-datasource.workspace = true common-error.workspace = true common-macro.workspace = true -common-query.workspace = true common-recordbatch.workspace = true common-telemetry.workspace = true common-test-util = { workspace = true, optional = true } diff --git a/src/flow/Cargo.toml b/src/flow/Cargo.toml index 4586e612693..fea7b1eeda2 100644 --- a/src/flow/Cargo.toml +++ b/src/flow/Cargo.toml @@ -11,18 +11,13 @@ workspace = true api.workspace = true arrow.workspace = true arrow-schema.workspace = true -async-recursion = "1.0" async-trait.workspace = true -bytes.workspace = true catalog.workspace = true chrono.workspace = true client.workspace = true common-base.workspace = true common-config.workspace = true -common-datasource.workspace = true -common-decimal.workspace = true common-error.workspace = true -common-frontend.workspace = true common-function.workspace = true common-grpc.workspace = true common-macro.workspace = true @@ -35,41 +30,30 @@ common-stat.workspace = true common-telemetry.workspace = true common-time.workspace = true common-version.workspace = true -config.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true datafusion-physical-expr.workspace = true -datafusion-substrait.workspace = true datatypes.workspace = true -enum_dispatch = "0.3" futures.workspace = true get-size2 = "0.1.2" greptime-proto.workspace = true hostname.workspace = true -http.workspace = true humantime-serde.workspace = true itertools.workspace = true lazy_static.workspace = true meta-client.workspace = true -nom = "7.1.3" -num-traits = "0.2" -num_cpus.workspace = true operator.workspace = true -partition.workspace = true prometheus.workspace = true -prost.workspace = true query.workspace = true rand.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true session.workspace = true -smallvec.workspace = true snafu.workspace = true sql.workspace = true store-api.workspace = true -strum.workspace = true substrait.workspace = true table.workspace = true tokio.workspace = true diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index c98510bd7d5..5b15ad0a93c 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -141,8 +141,9 @@ fn relation_desc_from_output( columns: &[ColumnSchema], lineage: &[Option], source_primary_key_indices: &[usize], + plan: &LogicalPlan, ) -> RelationDesc { - let keys = source_primary_key_indices + let mut keys = source_primary_key_indices .iter() .filter_map(|source_index| { lineage @@ -150,6 +151,14 @@ fn relation_desc_from_output( .position(|index| index == &Some(*source_index)) }) .collect_vec(); + // DISTINCT defines identity from every output tuple; the source timestamp participates + // separately as the relation time index, so it is excluded from the tag key. + if matches!(plan, LogicalPlan::Distinct(Distinct::All(_))) { + keys = columns + .iter() + .positions(|column| !column.is_time_index()) + .collect(); + } let time_index = columns.iter().position(ColumnSchema::is_time_index); RelationDesc { typ: crate::repr::RelationType { @@ -849,8 +858,12 @@ impl StreamingEngine { let source_primary_key_indices = source_meta.primary_key_indices; stateless::validate_source_scan(&flow_plan, source_table_id, &source_schema)?; let (inferred_schema, lineage) = output_column_schemas(&flow_plan, &source_schema)?; - let inferred_relation = - relation_desc_from_output(&inferred_schema, &lineage, &source_primary_key_indices); + let inferred_relation = relation_desc_from_output( + &inferred_schema, + &lineage, + &source_primary_key_indices, + &flow_plan, + ); let sink_exists = self.fetch_table_pk_schema(sink_table_name).await?.is_some(); if !sink_exists { validate_auto_column_names(&inferred_schema)?; diff --git a/src/flow/src/adapter/flownode_impl.rs b/src/flow/src/adapter/flownode_impl.rs index 62e8597dbb6..cf4b6fd2d34 100644 --- a/src/flow/src/adapter/flownode_impl.rs +++ b/src/flow/src/adapter/flownode_impl.rs @@ -826,7 +826,16 @@ impl common_meta::node_manager::Flownode for FlowDualEngine { let query_ctx = request .header .and_then(|h| h.query_context) - .map(|ctx| ctx.into()); + .map(session::context::QueryContext::try_from) + .transpose() + .map_err(|err| { + to_meta_err(snafu::location!())( + crate::error::InvalidQuerySnafu { + reason: err.to_string(), + } + .build(), + ) + })?; match request.body { Some(flow_request::Body::Create(CreateRequest { flow_id: Some(task_id), diff --git a/src/flow/src/adapter/tests.rs b/src/flow/src/adapter/tests.rs index 6bc02f60786..6a2bf5ed816 100644 --- a/src/flow/src/adapter/tests.rs +++ b/src/flow/src/adapter/tests.rs @@ -20,7 +20,7 @@ use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as use datafusion::catalog::MemTable; use datafusion::datasource::provider_as_source; use datafusion_common::TableReference; -use datafusion_expr::LogicalPlanBuilder; +use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, Schema, SchemaBuilder}; use store_api::storage::{ConcreteDataType, TableId}; use table::metadata::{TableInfo, TableInfoBuilder, TableMetaBuilder}; @@ -194,10 +194,133 @@ fn stateless_distinct_preserves_direct_column_lineage() { let (output, lineage) = super::output_column_schemas(&plan, &source).unwrap(); assert_eq!(output[0].name, "dis"); assert_eq!(lineage, vec![Some(0)]); - let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &plan); assert_eq!(relation.typ.keys[0].column_indices, vec![0]); } +#[test] +fn stateless_distinct_output_is_always_the_sink_key() { + let source = Arc::new(Schema::new(vec![ + ColumnSchema::new("v", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new("k", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new("other", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ])); + // All cases share one provider so the key assertions exercise only plan/output metadata. + let provider = Arc::new( + MemTable::try_new( + Arc::new(datafusion::arrow::datatypes::Schema::new(vec![ + Field::new("v", ArrowDataType::Int32, true), + Field::new("k", ArrowDataType::Int32, true), + Field::new("other", ArrowDataType::Int32, true), + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + ])), + vec![vec![]], + ) + .unwrap(), + ); + let make_plan = |expressions: Vec, distinct: bool| { + let builder = LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(provider.clone()), + None, + ) + .unwrap() + .project(expressions) + .unwrap(); + if distinct { + builder.distinct().unwrap().build().unwrap() + } else { + builder.build().unwrap() + } + }; + let assert_key = |plan: &LogicalPlan, source_keys: &[usize], expected: &[usize]| { + let (output, lineage) = super::output_column_schemas(plan, &source).unwrap(); + let relation = super::relation_desc_from_output(&output, &lineage, source_keys, plan); + assert_eq!( + relation + .typ + .keys + .iter() + .map(|key| key.column_indices.clone()) + .collect::>(), + if expected.is_empty() { + vec![] + } else { + vec![expected.to_vec()] + } + ); + relation + }; + + // No source PK: all distinct output fields form the key. + assert_key(&make_plan(vec![datafusion_expr::col("v")], true), &[], &[0]); + // A surviving subset of a composite source PK cannot identify projected DISTINCT tuples. + assert_key( + &make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("other")], + true, + ), + &[0, 1], + &[0, 1], + ); + // Even a retained source PK needs the additional selected fields in the DISTINCT key. + assert_key( + &make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("other")], + true, + ), + &[0], + &[0, 1], + ); + // Aliases and computed output expressions are included by output position. + let alias_plan = make_plan( + vec![ + datafusion_expr::col("v").alias("renamed"), + (datafusion_expr::col("v") + datafusion_expr::col("other")).alias("computed"), + ], + true, + ); + let (output, lineage) = super::output_column_schemas(&alias_plan, &source).unwrap(); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &alias_plan); + assert_eq!( + output + .iter() + .map(|column| column.name.as_str()) + .collect_vec(), + ["renamed", "computed"] + ); + assert_eq!(relation.typ.keys[0].column_indices, vec![0, 1]); + + // A selected time index remains the relation time index, not part of the tag key. + let relation = assert_key( + &make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("ts")], + true, + ), + &[0], + &[0], + ); + assert_eq!(relation.typ.time_index, Some(1)); + + // Non-DISTINCT plans retain only inherited source-key lineage (or no key if none survives). + let ordinary = make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("other")], + false, + ); + assert_key(&ordinary, &[0, 1], &[0]); + assert_key(&ordinary, &[], &[]); +} + #[test] fn stateless_normalizes_dictionary_output_type() { let field = Field::new_dictionary("host", ArrowDataType::UInt32, ArrowDataType::Utf8, true); @@ -220,7 +343,7 @@ fn stateless_normalizes_dictionary_output_type() { assert_eq!(output[0].data_type, ConcreteDataType::string_datatype()); assert_eq!(lineage, vec![Some(0)]); - let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &plan); assert_eq!(relation.typ.keys[0].column_indices, vec![0]); } diff --git a/src/flow/src/batching_mode/batching_execution.rs b/src/flow/src/batching_mode/batching_execution.rs index 0a699a227be..298522b8137 100644 --- a/src/flow/src/batching_mode/batching_execution.rs +++ b/src/flow/src/batching_mode/batching_execution.rs @@ -16,6 +16,7 @@ use std::sync::Arc; +use datafusion_expr::LogicalPlan; use query::QueryEngineRef; use table::TableRef; @@ -37,6 +38,11 @@ pub trait BatchingExecution: Send + Sync + 'static { max_window_cnt: Option, ) -> ExecuteOnceOutcome; + /// Rewrite the completed query plan after incremental merging and before execution. + fn rewrite_plan(&self, _task: &BatchingTask, plan: LogicalPlan) -> Result { + Ok(plan) + } + /// Retire this execution instance, rejecting new work and requesting that any /// retained local work stop. This is not an acknowledgement of remote quiescence. fn stop(&self) {} diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 42fcc52d6f7..1a03bcbe9d4 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -91,7 +91,8 @@ impl HandlerMutable { pub enum FrontendClient { Distributed { meta_client: Arc, - chnl_mgr: ChannelManager, + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, query: QueryOptions, batch_opts: BatchingModeOptions, }, @@ -136,17 +137,19 @@ impl FrontendClient { batch_opts: BatchingModeOptions, ) -> Result { common_telemetry::info!("Frontend client build without auth"); + let cfg = ChannelConfig::new() + .connect_timeout(batch_opts.grpc_conn_timeout) + .timeout(Some(batch_opts.query_timeout)); + let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone()) + .context(InvalidClientConfigSnafu)?; + // Keep separate TLS config handles and pools for query and control traffic. + let query_channel_manager = ChannelManager::with_config(cfg.clone(), tls_config.clone()); + let control_channel_manager = ChannelManager::with_config(cfg, tls_config); + Ok(Self::Distributed { meta_client, - chnl_mgr: { - let cfg = ChannelConfig::new() - .connect_timeout(batch_opts.grpc_conn_timeout) - .timeout(Some(batch_opts.query_timeout)); - - let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone()) - .context(InvalidClientConfigSnafu)?; - ChannelManager::with_config(cfg, tls_config) - }, + query_channel_manager, + control_channel_manager, query, batch_opts, }) @@ -227,7 +230,8 @@ impl FrontendClient { frontends: &[Peer], ) -> Result, Error> { let Self::Distributed { - chnl_mgr, + query_channel_manager, + control_channel_manager, batch_opts, .. } = self @@ -240,10 +244,15 @@ impl FrontendClient { .iter() .map(|peer| { let addr = peer.addr.clone(); - let chnl_mgr = chnl_mgr.clone(); + let query_channel_manager = query_channel_manager.clone(); + let control_channel_manager = control_channel_manager.clone(); async move { - let client = Client::with_manager_and_urls(chnl_mgr, vec![addr.clone()]); + let client = Client::with_query_and_control_managers( + query_channel_manager, + control_channel_manager, + vec![addr.clone()], + ); let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); match tokio::time::timeout(probe_timeout, database.sql("SELECT 1")).await { @@ -282,7 +291,8 @@ impl FrontendClient { ) -> Result { let Self::Distributed { meta_client: _, - chnl_mgr, + query_channel_manager, + control_channel_manager, query: _, batch_opts, } = self @@ -302,7 +312,11 @@ impl FrontendClient { for peer in frontends { let addr = peer.addr.clone(); - let client = Client::with_manager_and_urls(chnl_mgr.clone(), vec![addr.clone()]); + let client = Client::with_query_and_control_managers( + query_channel_manager.clone(), + control_channel_manager.clone(), + vec![addr.clone()], + ); let database = Database::new(catalog, schema, client); let db = DatabaseWithPeer::new(database, peer); match db.try_select_one().await { @@ -1260,7 +1274,11 @@ mod tests { let database = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls([addr.as_str()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.as_str()], + ), ); let db = DatabaseWithPeer::new( database, @@ -1297,7 +1315,11 @@ mod tests { let database = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls([addr.as_str()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.as_str()], + ), ); let db = DatabaseWithPeer::new( database, diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index b9710a25a0d..55307dbcdec 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -208,6 +208,25 @@ impl TaskState { self.pending_fenced_repair.as_ref() } + /// Start repairing explicit bounded windows under a frozen high `H`. + /// + /// The caller supplies the complete repair scope. Live dirty windows are + /// left unchanged so signals received after `H` remain separate. + pub fn start_fenced_repair_windows( + &mut self, + high: BTreeMap, + windows: Vec<(Timestamp, Timestamp)>, + ) { + let mut pending_windows = self.dirty_time_windows.clone(); + pending_windows.clean(); + pending_windows.add_windows(windows); + self.pending_fenced_repair = Some(FencedRepair { + high, + pending_windows, + }); + self.checkpoint_mode = CheckpointMode::FullSnapshot; + } + /// Finish the fenced repair and promote the frozen high watermark to the /// checkpoint map. Incremental-disabled flows stay in FullSnapshot mode. pub fn finish_fenced_repair(&mut self) -> Option> { @@ -259,22 +278,18 @@ impl TaskState { task_ctx: Option<&BatchingTask>, ) -> Result, Error> { if let Some(repair) = self.pending_fenced_repair.as_mut() { - let expr = repair.pending_windows.gen_filter_exprs( + // Fenced windows are an explicit frozen repair scope. They must not + // be pruned by the moving live-data expiration boundary, and an + // empty repair must remain active until its high watermark is + // explicitly finished. + return repair.pending_windows.gen_filter_exprs( col_name, - expire_lower_bound, + None, window_size, window_cnt, flow_id, task_ctx, - )?; - if expr.is_some() || !repair.pending_windows.is_empty() { - return Ok(expr); - } - - // All pending repair windows may have expired during merge. Clear - // the empty repair so this call can fall back to live dirty windows - // instead of routing future executions to an empty queue forever. - self.pending_fenced_repair = None; + ); } self.dirty_time_windows.gen_filter_exprs( @@ -1393,6 +1408,108 @@ mod test { assert_eq!(state.dirty_time_windows.len(), 2); } + #[test] + fn test_explicit_fenced_repair_windows_keep_live_windows_separate() { + let mut state = state_with_past_update(Duration::from_secs(1)); + state + .dirty_time_windows + .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(1_000))); + let high = BTreeMap::from([(1, 10)]); + state.start_fenced_repair_windows( + high.clone(), + vec![(Timestamp::new_second(10), Timestamp::new_second(15))], + ); + + assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot); + assert_eq!(state.pending_fenced_repair().unwrap().high(), &high); + assert_eq!( + state + .pending_fenced_repair() + .unwrap() + .pending_windows() + .len(), + 1 + ); + assert_eq!(state.dirty_time_windows.len(), 1); + + let filter = state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .unwrap(); + assert_eq!( + filter.time_ranges, + vec![(Timestamp::new_second(10), Timestamp::new_second(15))] + ); + state + .dirty_time_windows + .add_window(Timestamp::new_second(20), Some(Timestamp::new_second(25))); + state.restore_scoped_windows(&filter); + + assert_eq!( + state + .pending_fenced_repair() + .unwrap() + .pending_windows() + .len(), + 1 + ); + assert_eq!(state.dirty_time_windows.len(), 2); + } + + #[test] + fn test_explicit_fenced_repair_keeps_empty_high_and_old_windows() { + let mut state = state_with_past_update(Duration::from_secs(1)); + state + .dirty_time_windows + .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(1_000))); + let high = BTreeMap::from([(1, 10)]); + state.start_fenced_repair_windows(high.clone(), Vec::new()); + + assert!( + state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .is_none() + ); + assert_eq!(state.pending_fenced_repair().unwrap().high(), &high); + assert_eq!(state.dirty_time_windows.len(), 1); + assert_eq!(state.finish_fenced_repair(), Some(high)); + + state.start_fenced_repair_windows( + BTreeMap::from([(1, 11)]), + vec![(Timestamp::new_second(-20), Timestamp::new_second(-15))], + ); + let filter = state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .unwrap(); + assert_eq!( + filter.time_ranges, + vec![(Timestamp::new_second(-20), Timestamp::new_second(-15))] + ); + } + #[test] fn test_disable_incremental_persists_full_snapshot_mode() { let query_ctx = QueryContext::arc(); diff --git a/src/flow/src/batching_mode/table_creator.rs b/src/flow/src/batching_mode/table_creator.rs index 73a17f94967..778ee939a8e 100644 --- a/src/flow/src/batching_mode/table_creator.rs +++ b/src/flow/src/batching_mode/table_creator.rs @@ -254,15 +254,22 @@ mod test { use std::sync::Arc; use api::v1::column_def::try_as_column_schema; + use catalog::RegisterTableRequest; + use catalog::memory::new_memory_catalog_manager; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use datafusion::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit, }; use datafusion_common::DFSchema; use datafusion_expr::logical_plan::EmptyRelation; use datatypes::prelude::ConcreteDataType; - use datatypes::schema::ColumnSchema; + use datatypes::schema::{ColumnSchema, Schema}; use pretty_assertions::assert_eq; + use query::options::QueryOptions; + use query::{QueryEngineFactory, QueryEngineRef}; use session::context::QueryContext; + use table::metadata::{TableInfoBuilder, TableMetaBuilder}; + use table::test_util::EmptyTable; use super::*; use crate::adapter::{AUTO_CREATED_PLACEHOLDER_TS_COL, AUTO_CREATED_UPDATE_AT_TS_COL}; @@ -308,6 +315,120 @@ mod test { assert!(columns[2].is_time_index()); } + /// Creates a query engine holding a Prometheus shaped table `http_requests`: tags + /// (`host`, `idc`), a single f64 value column (`val`) and a time index (`ts`), so that + /// TQL queries can be planned against it. + fn create_tql_test_query_engine() -> QueryEngineRef { + let catalog_list = new_memory_catalog_manager().unwrap(); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("idc", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true), + ]))) + // `host` and `idc` are tags, `val` is the value column + .primary_key_indices(vec![1, 2]) + .value_indices(vec![3]) + .engine("mito".to_string()) + .next_column_id(1026) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name("http_requests".to_string()) + .meta(table_meta) + .build() + .unwrap(); + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "http_requests".to_string(), + table_id: 1026, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + + QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + } + + /// A TQL `count_values` flow must keep the generated label as a primary key of the auto + /// created sink table. `count_values("status_code", http_requests)` groups by the sample + /// value column and projects the label as a unary scalar expression of that column + /// (`prom_float_to_string(val) AS status_code`), so the label only derives from a group by + /// column instead of referencing it directly. + /// + /// The plan is built with the same path a flow task uses (including the DataFusion + /// optimizers, which may rewrite the shape of the alias), and the assertion covers both. + #[tokio::test] + async fn test_tql_count_values_generated_label_is_primary_key() { + let query_engine = create_tql_test_query_engine(); + let ctx = QueryContext::arc(); + + for optimize in [false, true] { + let plan = sql_to_df_plan( + ctx.clone(), + query_engine.clone(), + r#"TQL EVAL (0, 15, '5s') count_values("status_code", http_requests)"#, + optimize, + ) + .await + .unwrap(); + let plan_display = plan.display_indent_schema().to_string(); + let expr = create_table_with_expr( + &plan, + &[ + "greptime".to_string(), + "public".to_string(), + "sink".to_string(), + ], + &QueryType::Tql, + ) + .unwrap(); + let columns = expr + .column_defs + .iter() + .map(|column| try_as_column_schema(column).unwrap()) + .collect::>(); + + assert_eq!( + vec!["status_code".to_string()], + expr.primary_keys, + "optimize={optimize}, plan:\n{plan_display}" + ); + assert_eq!( + "ts", expr.time_index, + "optimize={optimize}, plan:\n{plan_display}" + ); + // the aggregation output is a value column, the generated label is a tag column + assert_eq!( + "count(http_requests.val)", columns[0].name, + "optimize={optimize}, plan:\n{plan_display}" + ); + assert_eq!(ConcreteDataType::float64_datatype(), columns[0].data_type); + assert_eq!("ts", columns[1].name); + assert!(columns[1].is_time_index()); + assert_eq!("status_code", columns[2].name); + assert_eq!(ConcreteDataType::string_datatype(), columns[2].data_type); + } + } + #[tokio::test] async fn test_gen_create_table_sql() { let query_engine = create_test_query_engine(); diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index c0ee6d1c1b5..c2c3082995a 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -18,21 +18,29 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use api::v1::{CreateTableExpr, TableName}; use catalog::CatalogManagerRef; +use catalog::kvbackend::KvBackendCatalogManager; use client::OutputWithMetrics; use common_error::ext::BoxedError; +use common_meta::key::schema_name::SchemaNameKey; +use common_query::OutputData; use common_query::logical_plan::breakup_insert_plan; use common_telemetry::tracing::warn; use common_telemetry::{debug, info}; -use common_time::Timestamp; +use common_time::{TimeToLive, Timestamp}; use datafusion::datasource::DefaultTableSource; use datafusion::sql::unparser::expr_to_sql; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::utils::quote_identifier; use datafusion_common::{DFSchemaRef, ScalarValue, TableReference}; -use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit}; +use datafusion_expr::{DmlStatement, LogicalPlan, Projection, WriteOp, col, lit}; use datatypes::schema::Schema; +use datatypes::vectors::Helper; +use futures::TryStreamExt; use query::QueryEngineRef; -use query::options::FLOW_INCREMENTAL_MODE; +use query::options::{ + FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE, + FLOW_RETURN_REGION_SEQ, +}; use query::query_engine::DefaultSerializer; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; @@ -55,7 +63,7 @@ use crate::batching_mode::state::{ use crate::batching_mode::table_creator::{QueryType, create_table_with_expr}; use crate::batching_mode::time_window::TimeWindowExpr; use crate::batching_mode::utils::{ - AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, + AddFilterRewriter, ColumnMatcherRewriter, analyze_incremental_aggregate_plan, df_plan_to_sql, gen_plan_with_matching_schema_and_values, get_table_info_df_schema, sql_to_df_plan, }; use crate::df_optimizer::apply_df_optimizer; @@ -192,6 +200,56 @@ fn encode_insert_plan_request( }) } +fn recovery_aggregate_input(plan: &LogicalPlan) -> Result { + let plan = match plan { + LogicalPlan::Projection(projection) => projection.input.as_ref(), + _ => plan, + }; + let LogicalPlan::Aggregate(aggregate) = plan else { + return UnexpectedSnafu { + reason: "Recovery timestamp projection did not find an aggregate".to_string(), + } + .fail(); + }; + Ok(aggregate.input.as_ref().clone()) +} + +fn capture_recovery_batch_windows( + batch: &common_recordbatch::RecordBatch, + time_window_expr: &TimeWindowExpr, + windows: &mut BTreeSet<(Timestamp, Timestamp)>, +) -> Result<(), Error> { + if batch.num_columns() != 1 { + return UnexpectedSnafu { + reason: format!( + "Recovery timestamp projection returned {} columns instead of one", + batch.num_columns() + ), + } + .fail(); + } + let values = Helper::try_into_vector(batch.column(0).clone()) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + for index in 0..values.len() { + let timestamp = values.get(index).as_timestamp().context(UnexpectedSnafu { + reason: "Recovery timestamp projection returned a null or non-timestamp value" + .to_string(), + })?; + let (start, end) = time_window_expr.eval(timestamp)?; + let window = ( + start.context(UnexpectedSnafu { + reason: "Recovery time-window expression returned no lower bound".to_string(), + })?, + end.context(UnexpectedSnafu { + reason: "Recovery time-window expression returned no upper bound".to_string(), + })?, + ); + windows.insert(window); + } + Ok(()) +} + fn format_insert_target_columns(plan: &LogicalPlan) -> String { plan.schema() .fields() @@ -483,6 +541,377 @@ impl BatchingTask { Arc::new(query_ctx.fork()) } + /// Returns the retention lower bound aligned to this task's time window. + pub fn recovery_retention_lower_bound(&self) -> Result, Error> { + let Some(expire_after) = self.config.expire_after else { + return Ok(None); + }; + let expire_after = u64::try_from(expire_after).map_err(|_| { + UnexpectedSnafu { + reason: format!( + "Flow {} has negative expire_after {expire_after}", + self.config.flow_id + ), + } + .build() + })?; + let now = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|err| { + UnexpectedSnafu { + reason: format!("Failed to read recovery wall clock: {err}"), + } + .build() + })? + .as_secs() as i64, + ); + let lower = now + .sub_duration(Duration::from_secs(expire_after)) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + self.config + .time_window_expr + .as_ref() + .context(UnexpectedSnafu { + reason: "Recovery expiry requires a time-window expression".to_string(), + })? + .eval(lower)? + .0 + .context(UnexpectedSnafu { + reason: "Recovery expiry time-window expression returned no lower bound" + .to_string(), + }) + .map(Some) + } + + /// Verifies that recovery can still read all source data required by its retained scope. + pub async fn validate_recovery_retention( + &self, + retention_lower: Option, + windows: &[(Timestamp, Timestamp)], + ) -> Result<(), Error> { + for name in &self.config.source_table_names { + let table = self + .config + .catalog_manager + .table(&name[0], &name[1], &name[2], None) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + .context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} is unavailable for recovery retention validation", + self.config.flow_id, + name.join(".") + ), + })?; + let ttl = if let Some(ttl) = table.table_info().meta.options.ttl { + ttl + } else { + let manager = self + .config + .catalog_manager + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: format!( + "Flow {} cannot resolve inherited TTL for source table {} during recovery", + self.config.flow_id, + name.join(".") + ), + })?; + manager + .table_metadata_manager_ref() + .schema_manager() + .get(SchemaNameKey::new(&name[0], &name[1])) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + .context(UnexpectedSnafu { + reason: format!( + "Flow {} schema {}.{} is unavailable for recovery retention validation", + self.config.flow_id, name[0], name[1] + ), + })? + .ttl + .map(Into::into) + .unwrap_or(TimeToLive::Forever) + }; + + match ttl { + TimeToLive::Forever => {} + TimeToLive::Instant => { + return UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has instant TTL and cannot be recovered", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + TimeToLive::Duration(ttl) => { + let expire_after = self.config.expire_after.context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has TTL {ttl:?}, but recovery has no expire_after", + self.config.flow_id, + name.join(".") + ), + })?; + let expire_after = u64::try_from(expire_after).map_err(|_| { + UnexpectedSnafu { + reason: format!( + "Flow {} has negative expire_after {expire_after} during recovery retention validation", + self.config.flow_id + ), + } + .build() + })?; + if ttl <= Duration::from_secs(expire_after) { + return UnexpectedSnafu { + reason: format!( + "Flow {} source table {} TTL {ttl:?} must exceed expire_after {expire_after}s for recovery", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + let mut oldest = retention_lower.context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has finite TTL {ttl:?}, but recovery has no retained lower bound", + self.config.flow_id, + name.join(".") + ), + })?; + for (start, _) in windows { + oldest = oldest.min(*start); + } + let cutoff = Timestamp::current_millis() + .sub_duration(ttl) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if oldest <= cutoff { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery requires source table {} data from {oldest:?}, older than TTL {ttl:?} cutoff {cutoff:?}", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + } + } + } + Ok(()) + } + + /// Discovers the source time windows touched by the exact sequence range `(C, H]`. + /// + /// This compatibility wrapper captures the retention bound before discovery. + pub async fn capture_recovery_windows( + &self, + engine: &QueryEngineRef, + frontend_client: &FrontendClient, + lower: &BTreeMap, + ) -> Result<(BTreeMap, Vec<(Timestamp, Timestamp)>), Error> { + let retention_lower = self.recovery_retention_lower_bound()?; + self.capture_recovery_windows_since(engine, frontend_client, lower, retention_lower) + .await + } + + /// Discovers windows under `lower` using a caller-frozen retention bound. + /// + /// The caller owns the execution guard and freezes this bound with its recovery scope. + /// This read-only method leaves task state untouched. Terminal proof must cover every + /// region in `lower`; subset or pruned proofs are rejected. + pub async fn capture_recovery_windows_since( + &self, + engine: &QueryEngineRef, + frontend_client: &FrontendClient, + lower: &BTreeMap, + retention_lower: Option, + ) -> Result<(BTreeMap, Vec<(Timestamp, Timestamp)>), Error> { + if lower.is_empty() { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires nonempty lower sequence bounds", + self.config.flow_id + ), + } + .fail(); + } + if !self.sequence_range_capable().await? { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires sequence_range-capable sources", + self.config.flow_id + ), + } + .fail(); + } + + let query_ctx = self.query_context_snapshot(); + let plan = sql_to_df_plan(query_ctx, engine.clone(), &self.config.query, false).await?; + let input = recovery_aggregate_input(&plan)?; + let Some(analysis) = analyze_incremental_aggregate_plan(&plan)? else { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires a supported aggregate input", + self.config.flow_id + ), + } + .fail(); + }; + if !analysis.unsupported_exprs.is_empty() { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture has unsupported aggregate expressions: {:?}", + self.config.flow_id, analysis.unsupported_exprs + ), + } + .fail(); + } + let time_window_expr = self + .config + .time_window_expr + .as_ref() + .context(UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires a time-window expression", + self.config.flow_id + ), + })?; + let input = if let Some(retention_lower) = retention_lower { + let mut add_filter = AddFilterRewriter::new( + col(&time_window_expr.column_name).gt_eq(lit(to_df_literal(retention_lower)?)), + ); + input + .rewrite(&mut add_filter) + .with_context(|_| DatafusionSnafu { + context: "Failed to apply recovery expire_after filter".to_string(), + })? + .data + } else { + input + }; + let timestamp_plan = LogicalPlan::Projection( + Projection::try_new(vec![col(&time_window_expr.column_name)], Arc::new(input)) + .context(DatafusionSnafu { + context: "Failed to project recovery source timestamps".to_string(), + })?, + ); + let catalog = &self.config.sink_table_name[0]; + let schema = &self.config.sink_table_name[1]; + let timestamp_plan = timestamp_plan + .clone() + .transform_down_with_subqueries(|p| { + if let LogicalPlan::TableScan(mut table_scan) = p { + let resolved = table_scan.table_name.resolve(catalog, schema); + table_scan.table_name = resolved.into(); + Ok(Transformed::yes(LogicalPlan::TableScan(table_scan))) + } else { + Ok(Transformed::no(p)) + } + }) + .with_context(|_| DatafusionSnafu { + context: format!( + "Failed to fix table ref in recovery timestamp plan, plan={:?}", + timestamp_plan + ), + })? + .data; + let message = DFLogicalSubstraitConvertor {} + .encode(×tamp_plan, DefaultSerializer) + .context(SubstraitEncodeLogicalPlanSnafu)?; + let request = api::v1::QueryRequest { + query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())), + }; + let lower_json = serde_json::to_string(lower).map_err(|err| { + UnexpectedSnafu { + reason: format!("Failed to serialize recovery lower sequence bounds: {err}"), + } + .build() + })?; + let extensions = [ + (FLOW_RETURN_REGION_SEQ, "true"), + (FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE), + (FLOW_INCREMENTAL_AFTER_SEQS, lower_json.as_str()), + ]; + let mut peer_desc = None; + let result = frontend_client + .query_with_terminal_metrics( + catalog, + schema, + request, + &extensions, + &HashMap::new(), + &mut peer_desc, + ) + .await?; + let mut aligned_windows = BTreeSet::new(); + let metrics = result.metrics.clone(); + match result.output.data { + OutputData::AffectedRows(_) => { + return UnexpectedSnafu { + reason: "Recovery timestamp projection unexpectedly returned affected rows" + .to_string(), + } + .fail(); + } + OutputData::RecordBatches(batches) => { + for batch in batches.iter() { + capture_recovery_batch_windows(batch, time_window_expr, &mut aligned_windows)?; + } + } + OutputData::Stream(mut stream) => { + while let Some(batch) = stream + .try_next() + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + { + capture_recovery_batch_windows(&batch, time_window_expr, &mut aligned_windows)?; + } + } + } + if !metrics.is_ready() { + return UnexpectedSnafu { + reason: "Recovery timestamp projection ended without terminal metrics".to_string(), + } + .fail(); + } + let participating = metrics.participating_regions().context(UnexpectedSnafu { + reason: "Recovery timestamp projection has no participating-region proof".to_string(), + })?; + let high = metrics.region_watermark_map().context(UnexpectedSnafu { + reason: "Recovery timestamp projection has no terminal watermark proof".to_string(), + })?; + if participating.len() != lower.len() + || high.len() != lower.len() + || participating + .iter() + .any(|region| !lower.contains_key(region)) + || lower + .iter() + .any(|(region, low)| high.get(region).is_none_or(|watermark| watermark < low)) + || high.keys().any(|region| !participating.contains(region)) + { + return UnexpectedSnafu { + reason: format!( + "Recovery timestamp projection returned incomplete or regressing terminal proof: lower={lower:?}, participating={participating:?}, high={high:?}" + ), + } + .fail(); + } + Ok(( + high.into_iter().collect(), + aligned_windows.into_iter().collect(), + )) + } + /// Validates that the sink table schema can accept this flow's ordinary output. pub async fn validate_sink_table_schema(&self, engine: &QueryEngineRef) -> Result<(), Error> { self.validate_sink_table_schema_with_values(engine, &BTreeMap::new()) @@ -828,6 +1257,10 @@ impl BatchingTask { return Ok(None); } let plan = incremental_plan.unwrap_or_else(|| plan.clone()); + let plan = match &self.execution { + Some(execution) => execution.rewrite_plan(self, plan)?, + None => plan, + }; let extensions = self .build_flow_query_extensions(incremental_safe, coverage.is_incremental_delta()) @@ -967,6 +1400,12 @@ impl BatchingTask { match res { Ok(res) => { let (affected_rows, _) = res.output.extract_rows_and_cost(); + if matches!(&res.output.data, common_query::OutputData::AffectedRows(_)) { + res.metrics.wait_ready().await; + } + if let Some(error) = res.metrics.completion_error() { + warn!("Flow {flow_id} completed with terminal metrics error: {error}"); + } debug!( "Flow {flow_id} executed, affected_rows: {affected_rows:?}, elapsed: {:?}, watermark: {:?}", elapsed, diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 7a56c70823d..29a15600f15 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::task::Poll; +use std::task::{Context, Poll}; +use std::time::Duration; use catalog::memory::MemoryCatalogManager; use catalog::{DeregisterTableRequest, RegisterTableRequest}; @@ -27,18 +29,21 @@ use common_error::status_code::StatusCode; use common_query::Output; use common_recordbatch::RecordBatch; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; +use common_time::{TimeToLive, Timestamp}; use datatypes::data_type::ConcreteDataType as CDT; -use datatypes::schema::ColumnSchema; +use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{ TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef, }; use pretty_assertions::assert_eq; +use prost::Message; use query::options::{ FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY, FLOW_SCHEDULED_TIME_MILLIS, FLOW_SINK_TABLE_ID, QueryOptions, }; use session::context::QueryContext; use snafu::ResultExt; +use substrait::substrait_proto_df::proto::Plan; use table::Table; use table::metadata::FilterPushDownType; use table::test_util::MemTable; @@ -143,6 +148,99 @@ impl crate::BatchingExecution for RetainingExecution { } } +/// Records the completed plan handed to it by the task. +#[derive(Default)] +struct PlanRecordingExecution { + plans: std::sync::Mutex>, +} + +#[async_trait::async_trait] +impl crate::BatchingExecution for PlanRecordingExecution { + async fn execute_once( + self: Arc, + _guard: BatchingExecutionGuard, + _task: &BatchingTask, + _engine: &QueryEngineRef, + _frontend: &Arc, + _max_window_cnt: Option, + ) -> ExecuteOnceOutcome { + ExecuteOnceOutcome { + new_query: None, + result: Ok(None), + } + } + + fn rewrite_plan(&self, _task: &BatchingTask, plan: LogicalPlan) -> crate::Result { + self.plans.lock().unwrap().push(plan.clone()); + Ok(plan) + } +} + +#[tokio::test] +async fn test_execution_hook_receives_the_completed_incremental_plan() { + let sink_table = "hook_sink"; + let query = "SELECT max(number) AS number, ts FROM numbers_with_ts GROUP BY ts"; + let TestTaskParts { + task, query_engine, .. + } = new_test_task_engine_and_plan_with_query(query, sink_table).await; + register_auto_created_aggregate_sink(&query_engine, sink_table); + + let ctx = task.state.read().unwrap().query_ctx.clone(); + let plan = sql_to_df_plan(ctx, query_engine.clone(), query, true) + .await + .unwrap(); + let (sink, _) = get_table_info_df_schema( + query_engine.engine_state().catalog_manager().clone(), + [ + "greptime".to_string(), + "public".to_string(), + sink_table.to_string(), + ], + ) + .await + .unwrap(); + let table_provider = Arc::new(DfTableProviderAdapter::new(sink)); + let table_source = Arc::new(DefaultTableSource::new(table_provider)); + let dml_plan = LogicalPlan::Dml(DmlStatement::new( + datafusion_common::TableReference::bare("test"), + table_source, + WriteOp::Insert(datafusion_expr::dml::InsertOp::Append), + Arc::new(plan), + )); + + task.state + .write() + .unwrap() + .advance_checkpoints(HashMap::from([(1_u64, 10_u64)])); + + let execution = Arc::new(PlanRecordingExecution::default()); + let task = task.with_execution(Some(execution.clone())); + let (frontend_client, _handler) = + FrontendClient::from_empty_grpc_handler(QueryOptions::default()); + let frontend = Arc::new(frontend_client); + let dirty_restore = DirtyRestore::Unscoped(dirty_range(10, 15)); + + let _ = task + .execute_plan_unlocked( + &query_engine, + &frontend, + &dml_plan, + &dirty_restore, + &QueryCoverage::IncrementalDelta, + ) + .await; + + let recorded = execution.plans.lock().unwrap(); + let plan = recorded + .first() + .expect("the execution hook must see the plan that is dispatched"); + let plan_text = plan.to_string(); + assert!( + plan_text.contains("Left Join"), + "the hook must receive the completed delta-sink merge plan, got:\n{plan_text}" + ); +} + #[tokio::test] async fn test_execution_delegate_dispatch_is_serialized() { let TestTaskParts { @@ -722,6 +820,47 @@ fn register_auto_created_aggregate_sink(query_engine: &QueryEngineRef, table_nam memory_catalog.register_table_sync(request).unwrap(); } +async fn configure_source_ttl(query_engine: &QueryEngineRef, ttl: Option) { + let catalog_manager = query_engine.engine_state().catalog_manager(); + let source = catalog_manager + .table( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "numbers_with_ts", + None, + ) + .await + .unwrap() + .unwrap(); + let mut info = (*source.table_info()).clone(); + info.meta.options.ttl = ttl; + let source = Arc::new(Table::new( + Arc::new(info), + FilterPushDownType::Unsupported, + source.data_source(), + )); + let memory_catalog = catalog_manager + .as_any() + .downcast_ref::() + .unwrap(); + memory_catalog + .deregister_table_sync(DeregisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "numbers_with_ts".to_string(), + }) + .unwrap(); + memory_catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "numbers_with_ts".to_string(), + table_id: source.table_info().table_id(), + table: source, + }) + .unwrap(); +} + async fn configure_source_capability( query_engine: &QueryEngineRef, engine: &str, @@ -778,6 +917,127 @@ async fn configure_source_capability( .unwrap(); } +#[tokio::test] +async fn test_validate_recovery_retention_explicit_source_ttls() { + struct Case { + name: &'static str, + ttl: Option, + expire_after: Option, + retention_age: Option, + window_age: Option, + should_pass: bool, + } + + let cases = [ + Case { + name: "finite TTL exceeding expire_after is admissible", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: true, + }, + Case { + name: "TTL equal to expire_after is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "TTL shorter than expire_after is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(59 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "aligned window older than TTL is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: Some(Duration::from_secs(3 * 60 * 60)), + should_pass: false, + }, + Case { + name: "old retention cutoff is rejected even without windows", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(3 * 60 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "explicit forever is admissible", + ttl: Some(TimeToLive::Forever), + expire_after: None, + retention_age: None, + window_age: None, + should_pass: true, + }, + Case { + name: "instant TTL is rejected", + ttl: Some(TimeToLive::Instant), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + ]; + + for case in cases { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "sink", + ) + .await; + Arc::get_mut(&mut task.config).unwrap().expire_after = case.expire_after; + configure_source_ttl(&query_engine, case.ttl).await; + + let now = Timestamp::current_millis(); + let retention_lower = case.retention_age.map(|age| now.sub_duration(age).unwrap()); + let windows = case + .window_age + .map(|age| { + let start = now.sub_duration(age).unwrap(); + vec![(start, now)] + }) + .unwrap_or_default(); + assert_eq!( + task.validate_recovery_retention(retention_lower, &windows) + .await + .is_ok(), + case.should_pass, + "{}", + case.name + ); + } +} + +#[tokio::test] +async fn test_validate_recovery_retention_rejects_unknown_inherited_ttl() { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_test_task_engine_and_plan_with_query("SELECT number, ts FROM numbers_with_ts", "sink") + .await; + Arc::get_mut(&mut task.config).unwrap().expire_after = Some(60 * 60); + configure_source_ttl(&query_engine, None).await; + + assert!( + task.validate_recovery_retention(Some(Timestamp::current_millis()), &[],) + .await + .is_err() + ); +} + fn dirty_marker() -> DirtyTimeWindows { let mut dirty = DirtyTimeWindows::default(); dirty.set_dirty(); @@ -925,6 +1185,313 @@ impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError } } +struct RecoveryCaptureHandler { + output: std::sync::Mutex>, + lower: String, + expire_after: Option, + expiry_lower: std::sync::Mutex>, +} + +struct RecoveryMetricsStream { + schema: Arc, + batches: Vec, + metrics: Option, + fail: bool, +} + +impl futures::Stream for RecoveryMetricsStream { + type Item = common_recordbatch::error::Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + if self.fail { + self.fail = false; + return Poll::Ready(Some(Err(common_recordbatch::error::Error::External { + source: BoxedError::new(MockError::new(StatusCode::Unexpected)), + location: snafu::Location::new(file!(), line!(), column!()), + }))); + } + Poll::Ready(self.batches.pop().map(Ok)) + } +} + +impl common_recordbatch::RecordBatchStream for RecoveryMetricsStream { + fn name(&self) -> &str { + "RecoveryMetricsStream" + } + + fn schema(&self) -> Arc { + self.schema.clone() + } + + fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> { + None + } + + fn metrics(&self) -> Option { + self.batches + .is_empty() + .then(|| self.metrics.clone()) + .flatten() + } +} + +fn field_selection(expr: &substrait::substrait_proto_df::proto::Expression) -> Option { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::field_reference::ReferenceType; + use substrait::substrait_proto_df::proto::expression::reference_segment::ReferenceType as Segment; + + let RexType::Selection(field) = expr.rex_type.as_ref()? else { + return None; + }; + let ReferenceType::DirectReference(reference) = field.reference_type.as_ref()? else { + return None; + }; + let Segment::StructField(field) = reference.reference_type.as_ref()? else { + return None; + }; + Some(field.field) +} + +fn timestamp_literal(expr: &substrait::substrait_proto_df::proto::Expression) -> Option { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::literal::LiteralType; + + let RexType::Literal(literal) = expr.rex_type.as_ref()? else { + return None; + }; + let LiteralType::PrecisionTimestamp(timestamp) = literal.literal_type.as_ref()? else { + return None; + }; + match timestamp.precision { + 0 => Some(Timestamp::new_second(timestamp.value)), + 3 => Some(Timestamp::new_millisecond(timestamp.value)), + 6 => Some(Timestamp::new_microsecond(timestamp.value)), + 9 => Some(Timestamp::new_nanosecond(timestamp.value)), + _ => None, + } +} + +fn assert_recovery_timestamp_plan(plan: &Plan, expire_after: Option) -> Option { + use substrait::substrait_proto_df::proto::plan_rel::RelType as PlanRelType; + use substrait::substrait_proto_df::proto::rel::RelType; + + let function_names = plan + .extensions + .iter() + .filter_map(|extension| match extension.mapping_type.as_ref()? { + substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::MappingType::ExtensionFunction(function) => { + Some((function.function_anchor, function.name.as_str())) + } + _ => None, + }) + .collect::>(); + let [root] = plan.relations.as_slice() else { + panic!("recovery capture must encode exactly one relation"); + }; + let Some(PlanRelType::Root(root)) = root.rel_type.as_ref() else { + panic!("recovery capture must encode a root relation"); + }; + assert_eq!(root.names, ["ts"]); + let Some(RelType::Project(project)) = root.input.as_ref().and_then(|rel| rel.rel_type.as_ref()) + else { + panic!("recovery capture must project the raw timestamp"); + }; + assert_eq!(project.expressions.len(), 1); + assert_eq!(field_selection(&project.expressions[0]), Some(1)); + + fn visit( + rel: &substrait::substrait_proto_df::proto::Rel, + function_names: &HashMap, + scans: &mut Vec>, + has_number_filter: &mut bool, + expiry_lower: &mut Option, + ) { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::literal::LiteralType; + use substrait::substrait_proto_df::proto::function_argument::ArgType; + use substrait::substrait_proto_df::proto::read_rel::ReadType; + use substrait::substrait_proto_df::proto::rel::RelType; + + match rel.rel_type.as_ref().expect("recovery relation type") { + RelType::Aggregate(_) | RelType::Write(_) | RelType::Ddl(_) | RelType::Update(_) => { + panic!("recovery capture must not aggregate or write") + } + RelType::Read(read) => match read.read_type.as_ref() { + Some(ReadType::NamedTable(table)) => scans.push(table.names.clone()), + _ => panic!("recovery capture must read a named source table"), + }, + RelType::Filter(filter) => { + let condition = filter + .condition + .as_ref() + .expect("recovery filter condition"); + let RexType::ScalarFunction(function) = condition + .rex_type + .as_ref() + .expect("recovery filter expression") + else { + panic!("recovery filter must be a scalar function"); + }; + let args = function + .arguments + .iter() + .map(|arg| match arg.arg_type.as_ref() { + Some(ArgType::Value(expr)) => expr, + _ => panic!("recovery filter must use value arguments"), + }) + .collect::>(); + match function_names.get(&function.function_reference) { + Some(&"equal") + if args.len() == 2 + && field_selection(args[0]) == Some(0) + && matches!( + args[1].rex_type.as_ref(), + Some(RexType::Literal(literal)) + if matches!(literal.literal_type, Some(LiteralType::I64(42))) + ) => + { + *has_number_filter = true + } + Some(&"gte") if args.len() == 2 && field_selection(args[0]) == Some(1) => { + *expiry_lower = timestamp_literal(args[1]); + } + _ => panic!("unexpected recovery filter"), + } + visit( + filter.input.as_ref().expect("recovery filter input"), + function_names, + scans, + has_number_filter, + expiry_lower, + ); + } + RelType::Project(project) => visit( + project.input.as_ref().expect("recovery project input"), + function_names, + scans, + has_number_filter, + expiry_lower, + ), + _ => panic!("unexpected recovery relation"), + } + } + + let mut scans = Vec::new(); + let mut has_number_filter = false; + let mut expiry_lower = None; + visit( + project.input.as_ref().expect("recovery projection input"), + &function_names, + &mut scans, + &mut has_number_filter, + &mut expiry_lower, + ); + assert_eq!( + scans, + vec![vec![ + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + "numbers_with_ts".to_string(), + ]] + ); + assert!(has_number_filter, "recovery capture lost WHERE number = 42"); + assert_eq!(expiry_lower.is_some(), expire_after.is_some()); + expiry_lower +} + +#[async_trait::async_trait] +impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError + for RecoveryCaptureHandler +{ + async fn do_query( + &self, + query: api::v1::greptime_request::Request, + ctx: QueryContextRef, + ) -> std::result::Result { + assert_eq!(ctx.extension(FLOW_INCREMENTAL_MODE), Some("sequence_range")); + assert_eq!( + ctx.extension(FLOW_INCREMENTAL_AFTER_SEQS), + Some(self.lower.as_str()) + ); + assert_eq!( + ctx.extension(query::options::FLOW_RETURN_REGION_SEQ), + Some("true") + ); + let api::v1::greptime_request::Request::Query(request) = query else { + panic!("recovery capture must issue a query request"); + }; + let Some(api::v1::query_request::Query::LogicalPlan(plan)) = request.query else { + panic!("recovery capture must issue a logical plan read"); + }; + let plan = Plan::decode(plan.as_slice()).unwrap(); + *self.expiry_lower.lock().unwrap() = + assert_recovery_timestamp_plan(&plan, self.expire_after); + Ok(self.output.lock().unwrap().take().unwrap()) + } +} + +fn recovery_timestamp_schema() -> Arc { + Arc::new(Schema::new(vec![ + ColumnSchema::new("ts", CDT::timestamp_millisecond_datatype(), true).with_time_index(true), + ])) +} + +fn recovery_stream_output( + batches: Vec, + watermarks: Option)>>, + fail: bool, +) -> Output { + Output::new_with_stream(Box::pin(RecoveryMetricsStream { + schema: recovery_timestamp_schema(), + batches, + metrics: watermarks.map(|watermarks| RecordBatchMetrics { + region_watermarks: watermarks + .into_iter() + .map(|(region_id, watermark)| RegionWatermarkEntry { + region_id, + watermark, + }) + .collect(), + ..Default::default() + }), + fail, + })) +} + +fn recovery_timestamp_batch(values: Vec>) -> RecordBatch { + RecordBatch::new( + recovery_timestamp_schema(), + vec![Arc::new(TimestampMillisecondVector::from(values)) as VectorRef], + ) + .unwrap() +} + +fn seed_recovery_state( + task: &BatchingTask, + lower: &BTreeMap, +) -> (BTreeMap, String) { + let mut state = task.state.write().unwrap(); + state.advance_checkpoints(lower.iter().map(|(region, seq)| (*region, *seq)).collect()); + state.dirty_time_windows.add_window( + Timestamp::new_millisecond(20_000), + Some(Timestamp::new_millisecond(25_000)), + ); + ( + state.checkpoints().clone(), + format!("{:?}", state.dirty_time_windows), + ) +} + +fn assert_recovery_state_unchanged( + task: &BatchingTask, + checkpoints: &BTreeMap, + dirty: &str, +) { + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoints(), checkpoints); + assert_eq!(format!("{:?}", state.dirty_time_windows), dirty); +} + struct CaptureScheduledNowHandler { expected_extension: String, captured_sql: Arc>>, @@ -1277,6 +1844,331 @@ fn output_with_region_watermarks( result } +#[tokio::test] +async fn test_capture_recovery_windows_decodes_raw_timestamp_projection_with_where_and_expiry() { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT max(src.number) AS output_value, date_bin(INTERVAL '5 second', src.ts) AS output_window FROM numbers_with_ts AS src WHERE src.number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let expire_after = 5; + Arc::get_mut(&mut task.config) + .expect("test task config should be uniquely owned") + .expire_after = Some(expire_after); + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: Some(expire_after), + expiry_lower: std::sync::Mutex::new(None), + }); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade( + &(handler.clone() + as Arc), + ), + QueryOptions::default(), + ); + let (before_checkpoints, before_dirty) = seed_recovery_state(&task, &lower); + let before = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64, + ); + + let (high, windows) = task + .capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .unwrap(); + let after = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64, + ); + + assert_eq!(high, lower); + assert!(windows.is_empty()); + let expiry_lower = handler.expiry_lower.lock().unwrap().unwrap(); + let time_window_expr = task.config.time_window_expr.as_ref().unwrap(); + // The helper and samples may straddle a window boundary. + let aligned_expiry = |now: Timestamp| { + time_window_expr + .eval( + now.sub_duration(Duration::from_secs(expire_after as u64)) + .unwrap(), + ) + .unwrap() + .0 + .unwrap() + }; + assert!(aligned_expiry(before) <= expiry_lower && expiry_lower <= aligned_expiry(after)); + assert_eq!( + time_window_expr.eval(expiry_lower).unwrap().0, + Some(expiry_lower) + ); + assert_recovery_state_unchanged(&task, &before_checkpoints, &before_dirty); +} + +#[tokio::test] +async fn test_capture_recovery_windows_since_uses_supplied_retention_lower() { + let TestTaskParts { + task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let retention_lower = Timestamp::new_second(-10); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: Some(1), + expiry_lower: std::sync::Mutex::new(None), + }); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade( + &(handler.clone() + as Arc), + ), + QueryOptions::default(), + ); + + let (high, windows) = task + .capture_recovery_windows_since(&query_engine, &frontend, &lower, Some(retention_lower)) + .await + .unwrap(); + + assert_eq!(high, lower); + assert!(windows.is_empty()); + assert_eq!(*handler.expiry_lower.lock().unwrap(), Some(retention_lower)); +} + +#[tokio::test] +async fn test_capture_recovery_windows_terminal_proof_cases_leave_state_unchanged() { + struct Case { + name: &'static str, + lower: BTreeMap, + watermarks: Option)>>, + succeeds: bool, + } + + let cases = vec![ + Case { + name: "empty_stream_complete_proof_h_equals_c", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(10))]), + succeeds: true, + }, + Case { + name: "missing_metrics", + lower: BTreeMap::from([(1, 10)]), + watermarks: None, + succeeds: false, + }, + Case { + name: "missing_c_region", + lower: BTreeMap::from([(1, 10), (2, 20)]), + watermarks: Some(vec![(1, Some(10))]), + succeeds: false, + }, + Case { + name: "unknown_none_watermark", + lower: BTreeMap::from([(1, 10), (2, 20)]), + watermarks: Some(vec![(1, Some(10)), (2, None)]), + succeeds: false, + }, + Case { + name: "unexpected_region", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(10)), (2, Some(20))]), + succeeds: false, + }, + Case { + name: "regressing_h", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(9))]), + succeeds: false, + }, + ]; + + for case in cases { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + case.watermarks, + false, + ))), + lower: serde_json::to_string(&case.lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade(&handler_dyn), + QueryOptions::default(), + ); + let (before_checkpoints, before_dirty) = seed_recovery_state(&task, &case.lower); + + let result = task + .capture_recovery_windows(&query_engine, &frontend, &case.lower) + .await; + assert_eq!(result.is_ok(), case.succeeds, "{}: {result:?}", case.name); + if case.succeeds { + assert_eq!( + result.unwrap(), + (case.lower.clone(), vec![]), + "{}", + case.name + ); + } + assert_recovery_state_unchanged(&task, &before_checkpoints, &before_dirty); + } +} + +#[tokio::test] +async fn test_capture_recovery_windows_streams_sorted_unique_windows() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![ + recovery_timestamp_batch(vec![Some(10_000), Some(-1), Some(1_000)]), + recovery_timestamp_batch(vec![Some(11_000), Some(1_000), Some(-5_000)]), + ], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + let (_, windows) = task + .capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .unwrap(); + assert_eq!( + windows, + vec![ + ( + Timestamp::new_millisecond(-5_000), + Timestamp::new_millisecond(0) + ), + ( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(5_000) + ), + ( + Timestamp::new_millisecond(10_000), + Timestamp::new_millisecond(15_000) + ), + ] + ); +} + +#[tokio::test] +async fn test_capture_recovery_windows_rejects_stream_error() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + true, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + assert!( + task.capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .is_err() + ); +} + +#[tokio::test] +async fn test_capture_recovery_windows_rejects_null_timestamp() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![recovery_timestamp_batch(vec![None])], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + assert!( + task.capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .is_err() + ); +} + #[test] fn test_apply_query_result_to_state_advances_full_snapshot_to_incremental() { let query_ctx = QueryContext::arc(); @@ -2234,7 +3126,7 @@ async fn test_full_snapshot_scoped_plan_marks_checkpoint_advance_safe_only_after } #[tokio::test] -async fn test_expired_empty_fenced_repair_generates_scoped_base_repair_plan() { +async fn test_expired_fenced_repair_uses_frozen_scope() { let TestTaskParts { mut task, query_engine, @@ -2272,11 +3164,18 @@ async fn test_expired_empty_fenced_repair_generates_scoped_base_repair_plan() { ) .await .unwrap() - .expect("expired empty repair should fall back to live dirty"); + .unwrap(); - assert!(matches!(plan.coverage, QueryCoverage::ScopedBaseRepair)); - assert!(plan.coverage.snapshot_seqs().is_empty()); - assert!(task.state.read().unwrap().pending_fenced_repair().is_none()); + // A fenced repair is a frozen explicit scope: expiry cannot discard its + // old window or make this query consume the live post-fence signal. + assert!(matches!( + plan.coverage, + QueryCoverage::FencedRepairChunk { .. } + )); + let state = task.state.read().unwrap(); + assert!(state.pending_fenced_repair().is_some()); + assert!(state.fenced_repair_pending_is_empty()); + assert_eq!(state.dirty_time_windows.len(), 1); } #[tokio::test] diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 0220f4ef48d..267fa170cdd 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -1164,6 +1164,16 @@ pub fn df_plan_to_sql(plan: &LogicalPlan) -> Result { Ok(sql.to_string()) } +/// Returns whether `expr` directly renames the group by expression `group_expr`. +/// +/// Renames are identified only by expression name. Derived unary expressions are not group +/// keys, even when they reference a single group key column. A PromQL `count_values("label", +/// metric)` plan lands here too: it groups by the formatted sample value +/// (`prom_float_to_string(value)`) and projects that group key column as the generated label. +fn is_alias_of_group_expr(group_expr: &Expr, expr: &Expr) -> DfResult { + Ok(group_expr.name_for_alias()? == expr.name_for_alias()?) +} + /// Helper to find the innermost group by expr in schema, return None if no group by expr #[derive(Debug, Clone, Default)] pub struct FindGroupByFinalName { @@ -1220,24 +1230,28 @@ impl TreeNodeVisitor<'_> for FindGroupByFinalName { Ok(TreeNodeRecursion::Continue) } - /// deal with projection when going up with group exprs + /// Applies only direct, name-matched aliases when propagating group expressions upward. fn f_up(&mut self, node: &Self::Node) -> datafusion_common::Result { if let LogicalPlan::Projection(projection) = node { + // Only direct aliases of group expressions rename group keys. Derived unary + // expressions remain ordinary projected columns. for expr in &projection.expr { let Some(group_exprs) = &mut self.group_exprs else { return Ok(TreeNodeRecursion::Continue); }; if let datafusion_expr::Expr::Alias(alias) = expr { // if a alias exist, replace with the new alias - let mut new_group_exprs = group_exprs.clone(); + let mut renamed_group_expr = None; for group_expr in group_exprs.iter() { - if group_expr.name_for_alias()? == alias.expr.name_for_alias()? { - new_group_exprs.remove(group_expr); - new_group_exprs.insert(expr.clone()); + if is_alias_of_group_expr(group_expr, alias.expr.as_ref())? { + renamed_group_expr = Some(group_expr.clone()); break; } } - *group_exprs = new_group_exprs; + if let Some(group_expr) = renamed_group_expr { + group_exprs.remove(&group_expr); + group_exprs.insert(expr.clone()); + } } } } diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 6315ae3d679..1af8880a7d4 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -22,7 +22,6 @@ use common_recordbatch::recordbatch::merge_record_batches; use common_recordbatch::{RecordBatch, util}; use common_time::Timestamp; use datafusion_common::ScalarValue; -use datafusion_common::tree_node::TreeNode as _; use datafusion_expr::{GroupingSet, TableScanBuilder}; use datatypes::arrow::array::{Array, AsArray}; use datatypes::arrow::datatypes::{Float64Type, Int64Type, UInt64Type}; @@ -1033,6 +1032,34 @@ async fn test_find_group_by_exprs() { } } +#[tokio::test] +async fn test_find_group_by_exprs_does_not_replace_group_key_with_derived_expr() { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + let sql = "SELECT host, lower(host) AS host_lc, ts, SUM(val) AS total \ + FROM (SELECT CAST(number AS STRING) AS host, ts, number AS val FROM numbers_with_ts) \ + GROUP BY host, ts"; + + for optimize in [false, true] { + let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, optimize) + .await + .unwrap(); + let plan_display = plan.display_indent_schema().to_string(); + let mut group_finder = FindGroupByFinalName::default(); + plan.visit(&mut group_finder).unwrap(); + let group_names = group_finder.get_group_expr_names().unwrap_or_default(); + + assert!( + group_names.contains("host"), + "optimize={optimize}, group keys {group_names:?} must keep host:\n{plan_display}" + ); + assert!( + !group_names.contains("host_lc"), + "optimize={optimize}, derived host_lc must not replace host:\n{plan_display}" + ); + } +} + #[tokio::test] async fn test_analyze_incremental_aggregate_plan() { let query_engine = create_test_query_engine(); @@ -1243,6 +1270,64 @@ async fn test_rewrite_incremental_aggregate_allows_alias_wrapped_scan() { assert_eq!(rewritten_fields, analysis.output_field_names); } +#[tokio::test] +async fn test_count_values_flow_plan_keeps_generated_label_as_group_key() { + // PromQL `count_values("label", metric)` groups by the *formatted* sample value + // (`prom_float_to_string(value)`), so such a flow carries a scalar UDF inside its GROUP + // BY instead of the raw value column: + // + // Projection: count(val), ts, prom_float_to_string(val) AS v + // Aggregate: groupBy=[[ts, prom_float_to_string(CAST(val AS Float64))]], aggr=[[count(val)]] + // + // Two things must keep working for that shape: the TQL flow transport (a Substrait + // encoded insert plan) and the group-key resolution that derives the sink schema, which + // is what keeps the generated label a primary key of the auto-created sink table. + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + for optimize in [false, true] { + let plan = sql_to_df_plan( + ctx.clone(), + query_engine.clone(), + r#"TQL EVAL (0, 15, '5s') count_values("v", numbers_with_ts)"#, + optimize, + ) + .await + .unwrap(); + let plan_display = plan.display_indent_schema().to_string(); + + DFLogicalSubstraitConvertor {} + .encode(&plan, DefaultSerializer) + .unwrap_or_else(|err| panic!("optimize={optimize}: {err}\n{plan_display}")); + + let mut group_finder = FindGroupByFinalName::default(); + plan.visit(&mut group_finder).unwrap(); + let group_names = group_finder.get_group_expr_names().unwrap_or_default(); + assert!( + group_names.contains("v"), + "optimize={optimize}, group keys {group_names:?} must keep the generated label:\n{plan_display}" + ); + assert!( + group_names.contains("ts"), + "optimize={optimize}, group keys {group_names:?}:\n{plan_display}" + ); + + // The flow plan keeps the generated label as a group key name, and the plan shape + // (PromQL extension nodes above the aggregate) stays outside the incremental + // aggregate rewrite: `prepare_plan_for_incremental` only rewrites `QueryType::Sql` + // flows, so a `count_values` flow keeps running as a full snapshot. + let analysis = analyze_incremental_aggregate_plan(&plan).unwrap().unwrap(); + assert!( + analysis.group_key_names.contains(&"v".to_string()), + "optimize={optimize}, analysis {analysis:?}:\n{plan_display}" + ); + assert!( + !analysis.unsupported_exprs.is_empty(), + "optimize={optimize}, a count_values flow plan must not become an incremental \ + aggregate rewrite candidate: {analysis:?}\n{plan_display}" + ); + } +} + #[tokio::test] async fn test_analyze_incremental_aggregate_plan_rejects_having_filter() { let sql = diff --git a/src/flow/src/lib.rs b/src/flow/src/lib.rs index e9ef90e09d6..36dc7c21818 100644 --- a/src/flow/src/lib.rs +++ b/src/flow/src/lib.rs @@ -48,7 +48,10 @@ pub use batching_mode::task::{ QueryCoverage, TaskArgs, }; pub use batching_mode::time_window::{TimeWindowExpr, find_time_window_expr}; -pub use batching_mode::utils::sql_to_df_plan; +pub use batching_mode::utils::{ + IncrementalAggregateAnalysis, IncrementalAggregateMergeColumn, IncrementalAggregateMergeOp, + analyze_incremental_aggregate_plan, sql_to_df_plan, +}; pub(crate) use engine::{CreateFlowArgs, FlowId, TableName}; pub use error::{Error, Result}; pub use server::{FlownodeBuilder, FlownodeInstance, FlownodeServer, FlownodeServiceBuilder}; diff --git a/src/flow/src/metrics.rs b/src/flow/src/metrics.rs index 00f93d47ab8..db524c669a3 100644 --- a/src/flow/src/metrics.rs +++ b/src/flow/src/metrics.rs @@ -20,14 +20,6 @@ use prometheus::*; lazy_static! { pub static ref METRIC_FLOW_TASK_COUNT: IntGauge = register_int_gauge!("greptime_flow_task_count", "flow task count").unwrap(); - pub static ref METRIC_FLOW_INPUT_BUF_SIZE: IntGauge = - register_int_gauge!("greptime_flow_input_buf_size", "flow input buf size").unwrap(); - pub static ref METRIC_FLOW_INSERT_ELAPSED: HistogramVec = register_histogram_vec!( - "greptime_flow_insert_elapsed", - "flow insert elapsed", - &["table_id"] - ) - .unwrap(); pub static ref METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME: HistogramVec = register_histogram_vec!( "greptime_flow_batching_engine_query_time_secs", "flow batching engine query time(seconds)", @@ -101,8 +93,6 @@ lazy_static! { &["flow_id", "mode"], ) .unwrap(); - pub static ref METRIC_FLOW_RUN_INTERVAL_MS: IntGauge = - register_int_gauge!("greptime_flow_run_interval_ms", "flow run interval in ms").unwrap(); pub static ref METRIC_FLOW_ROWS: IntCounterVec = register_int_counter_vec!( "greptime_flow_processed_rows", "Count of rows flowing through the system.", @@ -115,10 +105,4 @@ lazy_static! { &["type"] ) .unwrap(); - pub static ref METRIC_FLOW_ERRORS: IntCounterVec = register_int_counter_vec!( - "greptime_flow_errors", - "Count of errors in flow processing", - &["code"] - ) - .unwrap(); } diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index d552d73cf39..fdaa5dbcfe4 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -34,11 +34,17 @@ remote datanodes via `operator`/`client`. | `heartbeat` | `src/frontend/src/heartbeat.rs` | Heartbeat to metasrv; handles suspend / cache invalidation | | `service_config` | `src/frontend/src/service_config/` | Per-protocol option structs | +Shared frontend cache counters live in `src/common/frontend/src/metrics.rs`, +including caches owned by pipeline and protocol handlers. + ## Request lifecycles - **SQL query** (`instance.rs`): `do_query_inner` handles parsing, interceptors, permission checks, timeout/cancellation, and delegates planning/execution to `StatementExecutor`. Distributed scans enter through `region_query.rs`. +- **Logical-table batch DDL** (`instance.rs`): dedicated SQL handler → bounded + parse/interception → all-member permissions → operator CREATE preparation → + one logical-table batch procedure. Keep pre-submission validation side-effect free. - **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` → `check_permission` → `operator`'s `Inserter` (schema validation, optional auto-create, partition routing, meter admission) → local `RegionServer` @@ -51,6 +57,28 @@ remote datanodes via `operator`/`client`. - Internal gRPC listeners mark requests with `Channel::Internal` in middleware (`server.rs`), including requests handled by Enterprise Flight wrappers. +- **Flight bulk insert** (`instance/grpc.rs`): initializes on the first batch after + the lazy schema handshake; checks permissions and reconciles missing columns + through `Inserter` once per stream, then reuses the refreshed table. + +- **OTLP trace lookup tables** (`instance/otlp/trace_ingest.rs`): `Instance` owns + a shared cache of confirmed service/operation writes, which bypass batching. + `otlp.trace_aux_cache_size` sets its estimated byte budget at startup; zero disables it. + Async main writes retain completion results; a frontend task writes auxiliary + rows only for confirmed chunks before populating the cache. + Cold async trace requests reserve a shared auxiliary slot before main writes, + bounded separately by the table batcher's `max_inflight_requests`; confirmed + cache hits bypass this admission. The slot lasts through auxiliary completion. + These tables are ingestion-managed; manual mutation requires restarting + serving frontends. + +- **Logical-table batching** (`instance/logical_batcher.rs`): `Services` initializes + one shared batcher for opted-in HTTP Prom and nonlegacy OTLP metric-engine + writes. Downstream routers can initialize it when enabling replacement endpoints. + `Instance::handle_otlp_metric_row_inserts` shares OTLP eligibility and dispatch + between converters and falls back for incompatible tables. + The schema adapter holds a weak instance reference to avoid an ownership cycle. + - **Table batching** (`instance/builder.rs`): protocol entry points opt in through `QueryContext`. The primary inserter prepares eligible ordinary-table writes for `servers::batcher::table::TablePendingRowsBatcher`. A separate execution-only diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index b59e1e83a4a..92803933a8b 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -19,7 +19,6 @@ workspace = true [dependencies] api.workspace = true -arc-swap = "1.0" async-stream.workspace = true async-trait.workspace = true auth.workspace = true @@ -42,7 +41,6 @@ common-macro.workspace = true common-memory-manager.workspace = true common-meta.workspace = true common-options.workspace = true -common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true @@ -56,12 +54,11 @@ datafusion-expr.workspace = true datatypes.workspace = true futures.workspace = true hostname.workspace = true -humantime.workspace = true humantime-serde.workspace = true lazy_static.workspace = true log-query.workspace = true meta-client.workspace = true -num_cpus.workspace = true +moka = { workspace = true, features = ["sync"] } opentelemetry-proto.workspace = true opentelemetry-semantic-conventions = { version = "0.31", features = ["semconv_experimental"] } operator.workspace = true @@ -72,7 +69,6 @@ prometheus.workspace = true promql-parser.workspace = true prost.workspace = true query.workspace = true -rand.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true @@ -100,4 +96,3 @@ serde_json.workspace = true strfmt = "0.2" temp-env.workspace = true tower.workspace = true -uuid.workspace = true diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 86ba2cb9bfb..90e7abc4517 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -61,13 +61,14 @@ pub struct FrontendOptions { pub http: HttpOptions, pub grpc: GrpcOptions, /// The internal gRPC options for the frontend service. - /// it provide the same service as the public gRPC service, just only for internal use. + /// It serves the same services as the public one plus the internal handler. + /// CORS is always off on it. pub internal_grpc: Option, pub mysql: MysqlOptions, pub postgres: PostgresOptions, pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, - /// Shared experimental ordinary-table batching; independent of Prom batching. + /// Ordinary-table batching with independent logical-table controls. pub pending_rows_batcher: PendingRowsBatcherOptions, pub prom_store: PromStoreOptions, pub jaeger: JaegerOptions, @@ -131,6 +132,7 @@ impl Configurable for FrontendOptions { "meta_client.metasrv_addrs", "event_recorder.event_types", "pending_rows_batcher.protocols", + "pending_rows_batcher.logical_table.protocols", ]) } } @@ -205,6 +207,7 @@ mod tests { use futures::Stream; use meta_client::MetaClientRef; use meta_client::client::MetaClientBuilder; + use servers::batcher::BatchingProtocol; use servers::grpc::{FlightCompression, GRPC_SERVER}; use servers::http::HTTP_SERVER; use servers::http::result::greptime_result_v1::GreptimedbV1Response; @@ -225,6 +228,35 @@ mod tests { type GrpcStream = Pin> + Send + Sync + 'static>>; + #[test] + fn test_logical_batcher_protocols_from_env() { + temp_env::with_vars( + [ + ( + "FRONTEND_LOGICAL_TEST__PENDING_ROWS_BATCHER__PROTOCOLS", + Some("otlp,influxdb"), + ), + ( + "FRONTEND_LOGICAL_TEST__PENDING_ROWS_BATCHER__LOGICAL_TABLE__PROTOCOLS", + Some("prom,otlp"), + ), + ], + || { + let options = + FrontendOptions::load_layered_options(None, "FRONTEND_LOGICAL_TEST").unwrap(); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); + assert_eq!( + options + .pending_rows_batcher + .logical_table + .unwrap() + .protocols, + vec![BatchingProtocol::Prom, BatchingProtocol::Otlp] + ); + }, + ); + } + #[test] fn test_batcher_protocols_from_env() { temp_env::with_vars( @@ -236,11 +268,8 @@ mod tests { let options = FrontendOptions::load_layered_options(None, "FRONTEND_BATCHER_TEST").unwrap(); assert_eq!( - options.pending_rows_batcher.protocols, - vec![ - servers::http::BatchingProtocol::Influxdb, - servers::http::BatchingProtocol::HttpSql - ] + options.pending_rows_batcher.table.protocols, + vec![BatchingProtocol::Influxdb, BatchingProtocol::HttpSql] ); }, ); @@ -252,6 +281,7 @@ mod tests { assert!( !defaults .pending_rows_batcher + .table .pending_rows_batching_enabled() ); let options: FrontendOptions = toml::from_str( @@ -263,9 +293,14 @@ max_batch_rows = 25 "#, ) .unwrap(); - assert_eq!(options.pending_rows_batcher.max_batch_rows, 25); - assert_eq!(options.pending_rows_batcher.protocols.len(), 2); - assert!(options.pending_rows_batcher.pending_rows_batching_enabled()); + assert_eq!(options.pending_rows_batcher.table.max_batch_rows, 25); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); + assert!( + options + .pending_rows_batcher + .table + .pending_rows_batching_enabled() + ); let serialized = toml::to_string(&options).unwrap(); let parsed: FrontendOptions = toml::from_str(&serialized).unwrap(); assert_eq!(options.influxdb, parsed.influxdb); @@ -280,7 +315,6 @@ max_batch_rows = 25 let enabled: FrontendOptions = toml::from_str("experimental_metric_export = true").unwrap(); assert!(enabled.experimental_metric_export); let toml_string = toml::to_string(&opts).unwrap(); - assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); assert_eq!(parsed.otlp, opts.otlp); assert_eq!(parsed.influxdb, opts.influxdb); @@ -610,7 +644,11 @@ max_batch_rows = 25 expected: std::result::Result<&str, (StatusCode, &str)>, ) { let addr = frontend.server_handlers().addr(GRPC_SERVER).unwrap(); - let client = Client::with_urls([addr.to_string()]); + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.to_string()], + ); let client = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); let response = client.sql("SELECT 1").await; diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 5a99ccbf584..3ae5ae6bc5e 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -21,6 +21,7 @@ mod import_packed; mod influxdb; mod jaeger; mod log_handler; +mod logical_batcher; mod logs; mod opentsdb; mod otlp; @@ -31,7 +32,7 @@ mod region_query; use std::collections::HashSet; use std::pin::Pin; use std::sync::atomic::AtomicBool; -use std::sync::{Arc, atomic}; +use std::sync::{Arc, OnceLock, atomic}; use std::time::{Duration, SystemTime}; use async_stream::stream; @@ -47,8 +48,10 @@ use catalog::process_manager::{ use client::OutputData; use common_base::Plugins; use common_base::cancellation::CancellableFuture; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::{BoxedError, ErrorExt}; use common_event_recorder::EventRecorderRef; +use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::key::TableMetadataManagerRef; @@ -78,6 +81,7 @@ use query::metrics::OnDone; use query::parser::{PromQuery, QueryStatement}; use query::query_engine::DescribeResult; use query::query_engine::options::{QueryOptions, validate_catalog_and_schema}; +use servers::batcher::logical_table::LogicalTablePendingRowsBatcher; use servers::error::{ self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu, OtlpMetricModeIncompatibleSnafu, UnexpectedResultSnafu, @@ -130,6 +134,7 @@ pub struct Instance { query_engine: QueryEngineRef, plugins: Plugins, inserter: InserterRef, + logical_batcher: Arc>>>, deleter: DeleterRef, table_metadata_manager: TableMetadataManagerRef, event_recorder: EventRecorderRef, @@ -138,6 +143,8 @@ pub struct Instance { slow_query_options: SlowQueryOptions, influxdb_default_merge_mode: InfluxdbMergeMode, trace_ingest_chunk_size: usize, + trace_aux_cache: otlp::TraceAuxCache, + trace_aux_limiter: Option, otlp_resource_info: bool, suspend: Arc, @@ -145,7 +152,7 @@ pub struct Instance { // first layer key: db-string // key: direct input metric name // value: if runs in legacy mode - otlp_metrics_table_legacy_cache: DashMap>, + otlp_metrics_table_legacy_cache: Arc>>, } impl Instance { @@ -582,6 +589,12 @@ fn fast_legacy_check( .iter() .filter_map(|name| cache.get(name)) .collect::>(); + CACHE_HIT + .with_label_values(&["otlp_metrics_legacy"]) + .inc_by(hit_cache.len() as u64); + CACHE_MISS + .with_label_values(&["otlp_metrics_legacy"]) + .inc_by((names.len() - hit_cache.len()) as u64); if !hit_cache.is_empty() { let hit_legacy = hit_cache.iter().any(|en| *en.value()); let hit_prom = hit_cache.iter().any(|en| !*en.value()); @@ -794,6 +807,98 @@ impl Instance { query_interceptor.post_execute(Output { data, meta }, query_ctx) } + async fn batch_create_logical_tables( + &self, + query: &str, + query_ctx: QueryContextRef, + ) -> Result> { + use servers::query_handler::sql::{ + MAX_LOGICAL_TABLE_DDL_BYTES, MAX_LOGICAL_TABLE_DDL_STATEMENTS, + }; + ensure!(!self.is_suspended(), error::SuspendedSnafu); + ensure!( + query.len() <= MAX_LOGICAL_TABLE_DDL_BYTES, + InvalidSqlSnafu { + err_msg: "logical-table DDL exceeds 1 MiB" + } + ); + let interceptor = self.plugins.get::>(); + let interceptor = interceptor.as_ref(); + let query = interceptor.pre_parsing(query, query_ctx.clone())?; + ensure!( + query.len() <= MAX_LOGICAL_TABLE_DDL_BYTES, + InvalidSqlSnafu { + err_msg: "logical-table DDL exceeds 1 MiB" + } + ); + let statements = parse_stmt(query.as_ref(), query_ctx.sql_dialect()).map_err(|_| { + InvalidSqlSnafu { + err_msg: "invalid logical-table DDL", + } + .build() + })?; + ensure!( + !statements.is_empty() && statements.len() <= MAX_LOGICAL_TABLE_DDL_STATEMENTS, + InvalidSqlSnafu { + err_msg: "logical-table batch requires 1 to 128 statements" + } + ); + let original_statements = interceptor.map(|_| statements.clone()); + let statements = interceptor.post_parsing(statements, query_ctx.clone())?; + ensure!( + !statements.is_empty() && statements.len() <= MAX_LOGICAL_TABLE_DDL_STATEMENTS, + InvalidSqlSnafu { + err_msg: "logical-table batch requires 1 to 128 statements" + } + ); + let measure_rewritten_sql = || { + let mut remaining = MAX_LOGICAL_TABLE_DDL_BYTES - (statements.len() - 1); + statements.iter().all(|statement| { + matches!(statement, Statement::CreateTable(create) if create.consume_sql_size_budget(&mut remaining)) + }) + }; + ensure!( + original_statements + .as_ref() + .is_none_or(|original| original == &statements) + || measure_rewritten_sql(), + InvalidSqlSnafu { + err_msg: "invalid or oversized rewritten logical-table DDL" + } + ); + for statement in &statements { + ensure!( + matches!(statement, Statement::CreateTable(_)), + InvalidSqlSnafu { + err_msg: "batch requires explicit Metric logical CREATE TABLE statements" + } + ); + interceptor.pre_execute(Some(statement), None, query_ctx.clone())?; + } + for statement in &statements { + self.check_sql_permission(statement, &query_ctx).await?; + check_permission(self.plugins.clone(), statement, &query_ctx)?; + } + let count = statements.len(); + let creates = statements + .into_iter() + .filter_map(|statement| { + if let Statement::CreateTable(create) = statement { + Some(create) + } else { + None + } + }) + .collect(); + self.statement_executor + .batch_create_logical_tables(creates, query_ctx.clone()) + .await + .context(TableOperationSnafu)?; + (0..count) + .map(|_| interceptor.post_execute(Output::new_with_affected_rows(0), query_ctx.clone())) + .collect() + } + #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_query")] async fn do_query_inner(&self, query: &str, query_ctx: QueryContextRef) -> Vec> { if self.is_suspended() { @@ -1125,6 +1230,21 @@ impl Instance { #[async_trait] impl SqlQueryHandler for Instance { + fn supports_metric_batch_ddl(&self) -> bool { + true + } + + async fn create_logical_tables( + &self, + query: &str, + query_ctx: QueryContextRef, + ) -> server_error::Result> { + self.batch_create_logical_tables(query, query_ctx) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu) + } + async fn do_query( &self, query: &str, @@ -2775,12 +2895,12 @@ mod tests { ) } - struct ReleasedExportSource { + struct CancellableExportSource { schema: GtSchemaRef, channels: std::sync::Mutex, oneshot::Receiver<()>)>>, } - impl DataSource for ReleasedExportSource { + impl DataSource for CancellableExportSource { fn get_stream( &self, _request: ScanRequest, @@ -2800,15 +2920,15 @@ mod tests { } #[tokio::test] - async fn test_metric_export_http_timeout_drains_ordinary_writer() { + async fn test_metric_export_http_timeout_cancels_ordinary_source() { let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout"); let (started_tx, started_rx) = oneshot::channel(); - let (release_tx, release_rx) = oneshot::channel(); + let (mut release_tx, release_rx) = oneshot::channel(); let info = test_table_info(1024, "source").unwrap(); let source = Arc::new(Table::new( Arc::new(info.clone()), FilterPushDownType::Unsupported, - Arc::new(ReleasedExportSource { + Arc::new(CancellableExportSource { schema: info.meta.schema.clone(), channels: std::sync::Mutex::new(Some((started_tx, release_rx))), }), @@ -2846,9 +2966,9 @@ mod tests { let server_task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap(); }); + let uri = reqwest::Url::from_directory_path(destination.path()).unwrap(); let sql = format!( - "COPY DATABASE greptime.public TO '{}/' WITH (experimental_metric_export='true')", - destination.path().display() + "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')" ); let request = tokio::spawn(async move { reqwest::Client::builder() @@ -2867,21 +2987,10 @@ mod tests { .unwrap(); let response = request.await.unwrap(); assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT); - // A dropped ordinary stream would close this receiver before release. - release_tx.send(()).unwrap(); - let file = destination.path().join("source.parquet"); - tokio::time::timeout(Duration::from_secs(5), async { - loop { - if std::fs::read(&file) - .is_ok_and(|bytes| bytes.len() > 8 && bytes.ends_with(b"PAR1")) - { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .unwrap(); + tokio::time::timeout(Duration::from_secs(5), release_tx.closed()) + .await + .unwrap(); + assert!(!destination.path().join("source.parquet").exists()); server_task.abort(); } @@ -3334,12 +3443,16 @@ mod tests { } #[tokio::test] - async fn test_event_recorder_is_exposed() -> TestResult<()> { + async fn test_instance_shared_state() -> TestResult<()> { let instance = test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) .await?; let _event_recorder = instance.event_recorder(); + assert!(Arc::ptr_eq( + &instance.otlp_metrics_table_legacy_cache, + &instance.clone().otlp_metrics_table_legacy_cache, + )); Ok(()) } @@ -3470,11 +3583,20 @@ mod tests { })?; let ctx = test_query_ctx(1); let handler: PipelineHandlerRef = Arc::new(instance.clone()); + let hits = CACHE_HIT.with_label_values(&["pipeline_table"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_table"]); + let before = (hits.get(), misses.get()); handler .get_pipeline("pipeline", None, ctx.clone()) .await .unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1)); + handler + .get_pipeline("pipeline", None, ctx.clone()) + .await + .unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); assert_permission_denied( PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await, ); @@ -3793,11 +3915,15 @@ mod tests { #[test] fn test_fast_legacy_check_is_read_only() { + let hits = CACHE_HIT.with_label_values(&["otlp_metrics_legacy"]); + let misses = CACHE_MISS.with_label_values(&["otlp_metrics_legacy"]); + let before = (hits.get(), misses.get()); let cache = DashMap::new(); cache.insert("metric1".to_string(), true); let names = vec!["metric1".to_string(), "metric2".to_string()]; assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap()); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1)); assert!(!cache.contains_key("metric2")); cache_legacy_mode(&cache, &names, true).unwrap(); @@ -3809,6 +3935,9 @@ mod tests { cache_incompatible.insert("metric1".to_string(), true); cache_incompatible.insert("metric2".to_string(), false); assert!(fast_legacy_check(&cache_incompatible, &names).is_err()); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); + assert_eq!(None, fast_legacy_check(&cache, &[]).unwrap()); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); } #[test] @@ -4168,6 +4297,214 @@ mod tests { Ok(()) } + #[tokio::test] + async fn logical_ddl_validates_and_authorizes_rewritten_batch_before_submission() + -> TestResult<()> { + struct Rewrite(&'static str); + impl SqlQueryInterceptor for Rewrite { + type Error = Error; + fn pre_parsing<'a>( + &self, + query: &'a str, + _: QueryContextRef, + ) -> Result> { + Ok(if self.0 == "sql_bytes" { + std::borrow::Cow::Owned(format!("{query}{}", " ".repeat(1024 * 1024 + 1))) + } else { + std::borrow::Cow::Borrowed(query) + }) + } + fn pre_execute( + &self, + statement: Option<&Statement>, + _: Option<&LogicalPlan>, + _: QueryContextRef, + ) -> Result<()> { + ensure!( + !(self.0 == "pre_execute" + && matches!(statement, + Some(Statement::CreateTable(create)) if create.name.to_string() == "allowed2")), + InvalidSqlSnafu { + err_msg: "second member rejected by interceptor" + } + ); + Ok(()) + } + fn post_parsing( + &self, + mut statements: Vec, + _: QueryContextRef, + ) -> Result> { + match self.0 { + "target" => { + statements[1] = parse_one_sql( + "CREATE TABLE denied (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')", + ) + } + "count" => statements = (0..129).map(|i| parse_one_sql(&format!( + "CREATE TABLE allowed_{i} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')" + ))).collect(), + "type" => statements[1] = parse_one_sql("SELECT 1"), + "secret_bytes" | "array_bytes" | "small_option" => { + if let Statement::CreateTable(create) = &mut statements[1] { + let value = "x".repeat(if self.0 == "small_option" { 10 } else { 2 * 1024 * 1024 }); + if self.0 == "array_bytes" { + create.options.insert_options("non_scalar", vec![value.as_str()].into()); + } else { + create.options.insert("secret_access_key".to_string(), value); + } + } + } + "partition_shape" => { + if let Statement::CreateTable(create) = &mut statements[1] { + create.partitions = Some(sql::statements::create::Partitions { + column_list: vec![], + exprs: vec![sqlparser::ast::Expr::Identifier(sqlparser::ast::Ident::new("hidden"))], + }); + } + } + "vector_hidden" => { + statements[1] = parse_one_sql("CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX, v VECTOR(3)) ENGINE=metric WITH (on_physical_table='phy')"); + if let Statement::CreateTable(create) = &mut statements[1] { + create.columns[1].column_def.options.push(sqlparser::ast::ColumnOptionDef { + name: None, + option: sqlparser::ast::ColumnOption::Comment("private-value".repeat(200_000)), + }); + } + } + mode if mode.starts_with("fulltext_") || mode.starts_with("skipping_") + || mode.starts_with("inverted_") || mode.starts_with("vector_") => { + if let Statement::CreateTable(create) = &mut statements[1] { + let (kind, value_type) = mode.split_once('_').unwrap(); + let value = "x".repeat(if value_type == "small" { 10 } else { 2 * 1024 * 1024 }); + let extensions = &mut create.columns[1].extensions; + let options = match kind { + "fulltext" => &mut extensions.fulltext_index_options, + "skipping" => &mut extensions.skipping_index_options, + "inverted" => &mut extensions.inverted_index_options, + _ => &mut extensions.vector_options, + }.get_or_insert_default(); + if kind == "vector" { + options.insert("dim".to_string(), "3".to_string()); + } + if value_type == "array" { + options.insert_options("non_scalar", vec![value.as_str()].into()); + } else { + options.insert("secret_access_key".to_string(), value); + } + } + } + "physical_null" | "physical_array" => { + let value = if self.0 == "physical_null" { "NULL" } else { "['x']" }; + statements[1] = parse_one_sql(&format!( + "CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy', physical_metric_table={value})" + )); + } + "ast_bytes" => { + if let Statement::CreateTable(create) = &mut statements[1] { + create.name = + ObjectName(vec![sqlparser::ast::ObjectNamePart::Identifier( + sqlparser::ast::Ident::new("t".repeat(1024 * 1024)), + )]); + } + } + _ => {} + } + Ok(statements) + } + } + for mode in [ + "none", + "target", + "count", + "type", + "sql_bytes", + "ast_bytes", + "secret_bytes", + "array_bytes", + "small_option", + "physical_null", + "physical_array", + "fulltext_secret", + "fulltext_array", + "fulltext_small", + "skipping_secret", + "skipping_array", + "skipping_small", + "inverted_secret", + "inverted_array", + "inverted_small", + "vector_secret", + "vector_array", + "vector_small", + "vector_hidden", + "partition_shape", + "pre_execute", + "cross_catalog", + "valid", + ] { + let catalog = + catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + let executor = Arc::new(MockProcedureExecutor::new(catalog.clone())); + let plugins = Plugins::new(); + plugins.insert::(Arc::new(RejectUnresolvedPermissionChecker)); + plugins.insert::>(Arc::new(Rewrite(mode))); + plugins.insert(QueryOptions { + disallow_cross_catalog_query: mode == "cross_catalog", + }); + let instance = test_instance_with_catalog_manager( + catalog, + test_table(1025, "target")?, + plugins, + None, + executor.clone(), + ) + .await?; + let create = |name| { + format!( + "CREATE TABLE {name} (ts TIMESTAMP TIME INDEX, host STRING) ENGINE=metric WITH (on_physical_table='phy')" + ) + }; + let sql = format!( + "{};{}", + create("allowed"), + create(match mode { + "none" => "denied", + "cross_catalog" => "other.public.allowed2", + _ => "allowed2", + }) + ); + let result = + SqlQueryHandler::create_logical_tables(&instance, &sql, QueryContext::arc()).await; + let submits = matches!(mode, "valid" | "small_option") + || (mode.ends_with("_small") && !mode.starts_with("vector_")); + if matches!(mode, "none" | "target") { + let err = result.unwrap_err(); + assert_eq!( + err.status_code(), + StatusCode::PermissionDenied, + "{mode}: {err}" + ); + } else { + assert_eq!( + result.unwrap_err().status_code(), + match mode { + _ if submits => StatusCode::Unsupported, + "cross_catalog" => StatusCode::AccessDenied, + _ => StatusCode::InvalidArguments, + }, + "{mode}" + ); + } + assert_eq!( + executor.submitted.lock().unwrap().len(), + usize::from(submits), + "{mode}" + ); + } + Ok(()) + } + /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must /// succeed with 0 affected rows (no error, no DDL task submitted), instead of /// returning `TableNotFound`. diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 01e9ef034e3..032722499f9 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -22,6 +22,7 @@ use catalog::kvbackend::KvBackendCatalogManager; use catalog::process_manager::ProcessManagerRef; use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; +use common_batcher::request_limiter::RequestLimiter; use common_datasource::object_store::LocalFileAccess; use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef}; @@ -51,7 +52,7 @@ use pipeline::pipeline_operator::PipelineOperator; use query::QueryEngineFactory; use query::region_query::RegionQueryHandlerFactoryRef; use servers::batcher::table::TablePendingRowsBatcher; -use servers::http::BatchingProtocol; +use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled}; use snafu::{OptionExt, ResultExt}; use crate::error::{self, DataFusionSnafu, ExternalSnafu, Result}; @@ -60,8 +61,9 @@ use crate::frontend::FrontendOptions; use crate::heartbeat::frontend_peer_addr; use crate::instance::Instance; use crate::instance::entity_graph::EntityGraphProviderImpl; +use crate::instance::otlp::TraceAuxCache; use crate::instance::region_query::FrontendRegionQueryHandler; -use crate::service_config::PendingRowsBatcherOptions; +use crate::service_config::BatcherOptions; /// The frontend [`Instance`] builder. pub struct FrontendBuilder { @@ -237,31 +239,30 @@ impl FrontendBuilder { }; // The execution-only inserter owns no batchers, avoiding an Arc cycle. let bulk_inserter = Arc::new(create_inserter()); - let build_batcher = - |options: &PendingRowsBatcherOptions| -> Option> { - if !options.pending_rows_batching_enabled() - || (self.options.prom_store.with_metric_engine - && options - .protocols - .iter() - .all(|protocol| *protocol == BatchingProtocol::Prom)) - { - return None; - } - TablePendingRowsBatcher::try_new( - options.pending_rows_flush_interval, - options.max_batch_rows, - options.max_concurrent_flushes, - options.worker_channel_capacity, - options.max_inflight_requests, - options.flow_notification_queue_capacity, - bulk_inserter.clone(), - ) - .map(|batcher| batcher as Arc) - }; + let build_batcher = |options: &BatcherOptions| -> Option> { + if !options.pending_rows_batching_enabled() + || (self.options.prom_store.with_metric_engine + && options + .protocols + .iter() + .all(|protocol| *protocol == BatchingProtocol::Prom)) + { + return None; + } + TablePendingRowsBatcher::try_new( + options.pending_rows_flush_interval, + options.max_batch_rows, + options.max_concurrent_flushes, + options.worker_channel_capacity, + options.max_inflight_requests, + options.flow_notification_queue_capacity, + bulk_inserter.clone(), + ) + .map(|batcher| batcher as Arc) + }; let inserter = Arc::new( create_inserter() - .with_pending_rows_batcher(build_batcher(&self.options.pending_rows_batcher)), + .with_pending_rows_batcher(build_batcher(self.options.table_batcher_options())), ); let deleter = Arc::new(Deleter::new( self.catalog_manager.clone(), @@ -382,10 +383,21 @@ impl FrontendBuilder { self.options.event_recorder.ttl, )), self.options.event_recorder.event_types.clone(), + self.options.event_recorder.flush_interval, )); admin_event_recorder.install(&event_recorder); + let batcher_options = self.options.table_batcher_options(); + let trace_aux_limiter = if batcher_options.pending_rows_batching_enabled() + && !pending_rows_batch_sync_enabled() + { + RequestLimiter::try_new(batcher_options.max_inflight_requests) + } else { + None + }; + Ok(Instance { + logical_batcher: Default::default(), frontend_peer_addr, experimental_metric_export: self.options.experimental_metric_export, catalog_manager: self.catalog_manager, @@ -399,10 +411,12 @@ impl FrontendBuilder { event_recorder, slow_query_recorder, process_manager, - otlp_metrics_table_legacy_cache: DashMap::new(), + otlp_metrics_table_legacy_cache: Arc::new(DashMap::new()), slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, + trace_aux_cache: TraceAuxCache::new(self.options.otlp.trace_aux_cache_size.as_bytes()), + trace_aux_limiter, otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), }) diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 13898f71247..7dfd97cd69d 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -395,6 +395,7 @@ impl Instance { let catalog_manager = self.catalog_manager().clone(); let plugins = self.plugins.clone(); let inserter = self.inserter.clone(); + let statement_executor = self.statement_executor.clone(); let ctx = ctx.clone(); let mut table_ref: Option = None; let mut table_checked = false; @@ -425,25 +426,37 @@ impl Instance { .context(PermissionSnafu)?; // Resolve table reference - table_ref = Some( - catalog_manager - .table( - &table_name.catalog_name, - &table_name.schema_name, - &table_name.table_name, - None, - ) - .await - .context(CatalogSnafu)? - .with_context(|| TableNotFoundSnafu { - table_name: table_name.to_string(), - })?, - ); + let table = catalog_manager + .table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + None, + ) + .await + .context(CatalogSnafu)? + .with_context(|| TableNotFoundSnafu { + table_name: table_name.to_string(), + })?; // Check permissions for the table let interceptor_ref = plugins.get::>(); let interceptor = interceptor_ref.as_ref(); - interceptor.pre_bulk_insert(table_ref.clone().unwrap(), ctx.clone())?; + interceptor.pre_bulk_insert(table.clone(), ctx.clone())?; + + // The client sends the schema lazily after the stream handshake. + // Reconcile it once, before writing the first batch. + table_ref = Some( + inserter + .ensure_bulk_insert_schema( + table, + request.record_batch.schema().as_ref(), + &ctx, + statement_executor.as_ref(), + ) + .await + .context(TableOperationSnafu)?, + ); table_checked = true; } diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs new file mode 100644 index 00000000000..255766ac14d --- /dev/null +++ b/src/frontend/src/instance/logical_batcher.rs @@ -0,0 +1,104 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::{Arc, Weak}; + +use api::v1::ColumnSchema; +use async_trait::async_trait; +use servers::batcher::BatchingProtocol; +use servers::batcher::logical_table::{LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer}; +use servers::error::{BatcherChannelClosedSnafu, Result}; +use session::context::QueryContextRef; +use snafu::OptionExt; + +use crate::frontend::FrontendOptions; +use crate::instance::Instance; + +impl Instance { + /// Initializes the shared logical batcher for enabled ingestion endpoints. + /// Disabled endpoints leave it uninitialized so downstream routers can enable it. + pub fn init_logical_batcher(self: &Arc, options: &FrontendOptions) { + let options_batcher = options.logical_batcher_options(); + let enabled = options_batcher + .protocols + .iter() + .any(|protocol| match protocol { + BatchingProtocol::Prom => options.prom_store.enable, + BatchingProtocol::Otlp => options.otlp.enable, + _ => false, + }); + if !options.prom_store.with_metric_engine + || !enabled + || !options_batcher.pending_rows_batching_enabled() + { + return; + } + self.logical_batcher.get_or_init(|| { + LogicalTablePendingRowsBatcher::try_new( + self.partition_manager().clone(), + self.node_manager().clone(), + self.catalog_manager().clone(), + self.table_flownode_set_cache().clone(), + true, + Arc::new(LogicalTables(Arc::downgrade(self))), + options_batcher.pending_rows_flush_interval, + options_batcher.max_batch_rows, + options_batcher.max_concurrent_flushes, + options_batcher.worker_channel_capacity, + options_batcher.max_inflight_requests, + options_batcher.flow_notification_queue_capacity, + ) + }); + } + + pub(crate) fn logical_batcher(&self) -> Option<&Arc> { + self.logical_batcher.get().and_then(Option::as_ref) + } +} + +// The instance owns the shared batcher; schema preparation must not keep that +// owner alive through a reference cycle. +struct LogicalTables(Weak); + +#[async_trait] +impl PendingRowsSchemaAlterer for LogicalTables { + async fn create_tables_if_missing_batch( + &self, + catalog: &str, + schema: &str, + tables: &[(&str, &[ColumnSchema])], + with_metric_engine: bool, + ctx: QueryContextRef, + ) -> Result<()> { + self.0 + .upgrade() + .context(BatcherChannelClosedSnafu)? + .create_tables_if_missing_batch(catalog, schema, tables, with_metric_engine, ctx) + .await + } + + async fn add_missing_prom_tag_columns_batch( + &self, + catalog: &str, + schema: &str, + tables: &[(&str, &[String])], + ctx: QueryContextRef, + ) -> Result<()> { + self.0 + .upgrade() + .context(BatcherChannelClosedSnafu)? + .add_missing_prom_tag_columns_batch(catalog, schema, tables, ctx) + .await + } +} diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index fa40c7eb59a..e4f395db457 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -18,6 +18,7 @@ pub mod trace_types; use std::sync::Arc; +use api::v1::RowInsertRequests; use async_trait::async_trait; use auth::{ OTLP_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget, @@ -27,10 +28,11 @@ use client::Output; use common_catalog::consts::{trace_operations_table_name, trace_services_table_name}; use common_error::ext::BoxedError; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_query::{OutputData, OutputMeta}; use common_telemetry::{tracing, warn}; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; -use operator::insert::{admit_row_insert_batches, admit_write}; +use operator::insert::{Inserter, admit_row_insert_batches, admit_write}; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use pipeline::{GreptimePipelineParams, PipelineWay}; use servers::error::{self, AuthSnafu, Result as ServerResult}; @@ -42,6 +44,7 @@ use servers::query_handler::{ MetricsIngestOutcome, OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, }; use session::context::QueryContextRef; +use session::protocol_ctx::OtlpMetricCtx; use snafu::ResultExt; use table::requests::{ OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM, SEMANTIC_PER_TABLE_INDEX_KEY, @@ -49,6 +52,7 @@ use table::requests::{ SOURCE_OPENTELEMETRY, }; +pub(crate) use self::trace_ingest::TraceAuxCache; use self::trace_ingest::trace_conventions; use crate::instance::Instance; use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; @@ -165,27 +169,14 @@ impl OpenTelemetryProtocolHandler for Instance { Arc::new(c) }; - // OTLP tables have one sample field in both the legacy and physical paths. - let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx.clone(), false, true) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - } else { - let physical_table = ctx - .extension(PHYSICAL_TABLE_PARAM) - .unwrap_or(GREPTIME_PHYSICAL_TABLE) - .to_string(); - self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - }?; + let output = self + .handle_otlp_metric_row_inserts(requests, ctx.clone(), &metric_ctx) + .await?; outcome.write_cost = output.meta.cost; - // Derived enrichment, written after the metric data is committed: - // failing here would make the client retry data the server already - // accepted, so every failure degrades to a warning instead. + // Derived enrichment follows the accepted metric submission, which may + // still be queued in asynchronous mode. Failures remain warning-only + // to avoid retrying metric data the server already accepted. if let Some(resource_info) = resource_info { let written = match self.check_row_insert_permission( &resource_info, @@ -326,6 +317,69 @@ impl OpenTelemetryProtocolHandler for Instance { } } +impl Instance { + /// Inserts converted OTLP metrics through the eligible batching or ordinary write path. + /// The caller must check table permissions and admit the complete request first. + pub async fn handle_otlp_metric_row_inserts( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + metric_ctx: &OtlpMetricCtx, + ) -> ServerResult { + let physical_table = ctx + .extension(PHYSICAL_TABLE_PARAM) + .unwrap_or(GREPTIME_PHYSICAL_TABLE) + .to_string(); + // The bulk path converts each request's time index unit to the + // destination table's unit during batch alignment, so no pre-gate + // alignment is needed here. + let batcher = self.logical_batcher().filter(|_| { + ctx.logical_batching_enabled() && !metric_ctx.is_legacy && metric_ctx.with_metric_engine + }); + let batcher = if batcher.is_some() + && self + .inserter + .can_batch_metric_rows(&requests, &ctx, &physical_table) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)? + { + batcher + } else { + None + }; + + // OTLP tables have one sample field in both the legacy and physical paths. + if let Some(batcher) = batcher { + let (rows, cost) = batcher + .submit_with(requests, ctx.clone(), |mut requests| { + let ctx = ctx.clone(); + async move { + Inserter::meter_row_inserts(&mut requests, &ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + }) + .await?; + Ok(Output::new( + OutputData::AffectedRows(rows as usize), + OutputMeta::new_with_cost(cost as _), + )) + } else if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { + self.handle_row_inserts(requests, ctx.clone(), false, true) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } else { + self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md index c123f1caae1..e308ebec184 100644 --- a/src/frontend/src/instance/otlp/README.md +++ b/src/frontend/src/instance/otlp/README.md @@ -175,6 +175,58 @@ Accounting follows the main-table write: accepted or rejected span counts; and - failure details are bounded before they are folded into `TraceIngestOutcome`. +A cache shared by each frontend's `Instance` clones skips previously confirmed +auxiliary writes for v0, v1, and v2. `otlp.trace_aux_cache_size` sets an estimated +memory budget across all catalogs, schemas, and trace tables on each frontend. +By default, the budget is 1/128 of the pod's cgroup memory limit or total host +memory, with a minimum of `32MiB`. If memory detection is unavailable, it uses +`32MiB`. An explicit size overrides this default; `0` disables caching. +The estimate includes key structs and allocated string capacity, charging shared +table names per entry. It excludes cache and allocator overhead, so it is not a +strict process-memory limit. Changes take effect on frontend or standalone +restart. There is no time-based expiry. +Services and operations are cached independently. Only a successful +auxiliary write populates the cache. Auxiliary writes always bypass batching, +including when main spans use asynchronous batching. Failures, eviction, +frontend restarts, and concurrent misses can cause repeat writes. Write cost +reflects only the writes actually performed. + +`greptime_frontend_cache_hit{type="otlp_trace_aux"}` and +`greptime_frontend_cache_miss{type="otlp_trace_aux"}` count distinct service and +operation keys at the first cache check in each request. +Requests with auxiliary admission count before acquiring a slot; other requests +count when preparing auxiliary writes. Admission and write rechecks do not count. +When the cache is disabled, all checked keys count as misses. + +Successful auxiliary writes with `skip_wal` enabled are cached too. If a +datanode crashes before flushing those rows while the frontend survives, the +cache can suppress their recreation, even by later WAL-enabled requests, until +eviction or frontend restart. Missing service/operation discovery rows in this +case are an accepted limitation; the cache does not guarantee durability beyond +the write's WAL policy. + +With asynchronous main-table batching, the response still acknowledges queue +admission. A frontend task waits for the storage results and writes auxiliary +rows only for successful chunks. Failed chunks do not populate the auxiliary +tables or cache. Deferred auxiliary failures are logged; they cannot change an +already returned response. Deferred writes retain the original request's row +admission, and their write cost is not included in the early response. + +Async trace requests with cache misses acquire one shared auxiliary-work slot +before submitting main chunks. Each frontend limits these requests separately +from main-table batching, using `pending_rows_batcher.max_inflight_requests`. +The slot is held through auxiliary completion, including failures. Requests +whose service/operation entries are all cached bypass this admission and need +no deferred task. Cache lookups and updates hold no lock across storage writes. + +The auxiliary tables are ingestion-managed. When caching is enabled, manual +`DROP`, `TRUNCATE`, `DELETE`, or other mutations require clearing the affected +caches or restarting all serving frontends before relying on ingestion to +populate the lookup rows again. +There is no automatic invalidation or public cache-clear command. In v0, fewer +duplicate rows are appended; distinct service/operation lookup results remain +the same. + A rejection detail carries the failing cause, not just a status code, so an unusable attribute value names its column, source value, and target type. Repeats of the same `(site, cause)` collapse into one entry with an occurrence count, and diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index e5b16c7d191..8b90a19ba49 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{BTreeMap, HashMap, HashSet}; +use std::mem; use std::sync::Arc; use api::helper::ColumnDataTypeWrapper; @@ -22,13 +23,17 @@ use api::v1::{ RowInsertRequest, RowInsertRequests, Rows, Value, }; use client::Output; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_meta::rpc::ddl::TriggerReason; +use common_query::WriteCompletion; use common_telemetry::{debug, warn}; use datatypes::prelude::ConcreteDataType; +use moka::sync::Cache; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use pipeline::PipelineWay; +use prometheus::IntCounter; use servers::error::{self, Result as ServerResult}; use servers::otlp; use servers::otlp::coerce::{coerce_value_data, is_supported_trace_coercion, trace_value_datatype}; @@ -43,6 +48,8 @@ use table::requests::{ SEMANTIC_TRACE_CONVENTIONS, SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2, }; +use table::table_name::TableName; +use tokio::sync::OwnedSemaphorePermit; use crate::instance::Instance; use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; @@ -52,7 +59,131 @@ use crate::instance::otlp::trace_types::{ is_trace_reconcile_candidate_type, prepare_trace_column_rewrites, push_observed_trace_type, truncate_for_diagnostics, }; -use crate::metrics::{OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS}; +use crate::metrics::{ + OTLP_TRACE_AUX_CACHE_HIT, OTLP_TRACE_AUX_CACHE_MISS, OTLP_TRACES_FAILURE_COUNT, + OTLP_TRACES_ROWS, +}; + +/// Confirmed auxiliary writes, shared across requests and instance clones. +/// +/// ponytail: these tables are ingestion-managed; add invalidation if external +/// mutations must be supported without restarting every serving frontend. +#[derive(Clone)] +pub(crate) struct TraceAuxCache { + entries: Cache, +} + +#[derive(Clone, Hash, PartialEq, Eq)] +struct TraceAuxCacheKey { + table: Arc, + entry: TraceAuxEntry, +} + +impl TraceAuxCacheKey { + /// Estimated key memory, excluding cache and allocator overhead. + fn size(&self) -> usize { + let entry_size = match &self.entry { + TraceAuxEntry::Service(service) => service.capacity(), + TraceAuxEntry::Operation(service, span, kind) => { + service.capacity() + span.capacity() + kind.capacity() + } + }; + // Charge the shared table name per entry to keep the estimate conservative. + mem::size_of::() + + mem::size_of::() + + self.table.catalog_name.capacity() + + self.table.schema_name.capacity() + + self.table.table_name.capacity() + + entry_size + } +} + +#[derive(Clone, Hash, PartialEq, Eq)] +enum TraceAuxEntry { + Service(String), + Operation(String, String, String), +} + +impl TraceAuxCache { + /// Creates a cache with an estimated byte budget; zero disables caching. + pub(crate) fn new(max_size: u64) -> Self { + Self { + entries: Cache::builder() + .max_capacity(max_size) + .weigher(|key: &TraceAuxCacheKey, _: &()| { + u32::try_from(key.size()).unwrap_or(u32::MAX) + }) + .build(), + } + } + + /// Removes confirmed writes and returns keys to remember after insertion. + /// Only the first cache check in a request supplies metrics. + fn filter( + &self, + table: TableName, + aux_data: &mut TraceAuxData, + metrics: Option<(&IntCounter, &IntCounter)>, + ) -> Vec { + let lookups = aux_data.services.len() + aux_data.operations.len(); + let table = Arc::new(table); + let mut pending = Vec::new(); + let mut retain = |entry| { + let key = TraceAuxCacheKey { + table: table.clone(), + entry, + }; + if self.entries.get(&key).is_some() { + false + } else { + pending.push(key); + true + } + }; + aux_data + .services + .retain(|service| retain(TraceAuxEntry::Service(service.clone()))); + aux_data.operations.retain(|(service, span, kind)| { + retain(TraceAuxEntry::Operation( + service.clone(), + span.clone(), + kind.clone(), + )) + }); + if let Some((hits, misses)) = metrics { + hits.inc_by((lookups - pending.len()) as u64); + misses.inc_by(pending.len() as u64); + } + pending + } + + fn record(&self, keys: Vec) { + for key in keys { + self.entries.insert(key, ()); + } + } + + /// Reserves background capacity only for misses, without holding a cache lock. + async fn acquire_for_misses( + &self, + limiter: &RequestLimiter, + table: TableName, + mut aux_data: TraceAuxData, + metrics: (&IntCounter, &IntCounter), + ) -> ServerResult>> { + self.filter(table.clone(), &mut aux_data, Some(metrics)); + if aux_data.is_empty() { + return Ok(None); + } + let permit = limiter + .acquire() + .await + .map_err(|_| error::BatcherChannelClosedSnafu.build())?; + // Another request may have confirmed these entries while admission waited. + self.filter(table, &mut aux_data, None); + Ok((!aux_data.is_empty()).then_some(permit)) + } +} /// Merge converted V2 rows only when their fixed schemas agree. fn merge_trace_v2_rows(rows: &mut Rows, batch: Rows) -> ServerResult<()> { @@ -101,11 +232,52 @@ struct TraceChunkIngestContext<'a> { /// Accumulates trace outcomes, auxiliary rows, and bounded failure details. struct TraceIngestState { + needs_aux_write: bool, aux_data: TraceAuxData, + pending_aux_data: Vec<(TraceAuxData, Vec)>, outcome: TraceIngestOutcome, failure_messages: TraceFailureMessages, } +impl TraceIngestState { + fn add_aux_data(&mut self, aux_data: TraceAuxData, completions: Vec) { + if completions.is_empty() { + self.aux_data.services.extend(aux_data.services); + self.aux_data.operations.extend(aux_data.operations); + } else if !aux_data.is_empty() { + self.pending_aux_data.push((aux_data, completions)); + } + } +} + +/// Keeps confirmed chunks even if another chunk fails or loses its completion sender. +async fn confirmed_trace_aux_data( + pending: Vec<(TraceAuxData, Vec)>, +) -> TraceAuxData { + let results = futures::future::join_all(pending.into_iter().map( + |(aux_data, completions)| async move { + for completion in completions { + completion.wait().await?; + } + Ok::<_, BoxedError>(aux_data) + }, + )) + .await; + let mut confirmed = TraceAuxData::default(); + for result in results { + match result { + Ok(aux_data) => { + confirmed.services.extend(aux_data.services); + confirmed.operations.extend(aux_data.operations); + } + Err(err) => { + warn!(err; "Skipping auxiliary trace rows after asynchronous main write failure") + } + } + } + confirmed +} + /// Bounded, deduplicated failure details for one trace request. /// /// Occurrences past [`TRACE_FAILURE_MESSAGE_LIMIT`] distinct failures are @@ -876,6 +1048,32 @@ impl Instance { conventions: &str, ctx: QueryContextRef, ) -> ServerResult { + let aux_limiter = self + .trace_aux_limiter + .as_ref() + .filter(|_| ctx.batching_enabled()); + // Reserve once per request, before any main chunks are submitted. Main + // batcher permits remain independent so multi-chunk requests can progress. + let aux_permit = if let Some(limiter) = aux_limiter { + let mut aux_data = TraceAuxData::default(); + for span in groups.iter().flat_map(|group| &group.spans) { + aux_data.observe_span(span); + } + self.trace_aux_cache + .acquire_for_misses( + limiter, + TableName::new(ctx.current_catalog(), ctx.current_schema(), &table_name), + aux_data, + (&OTLP_TRACE_AUX_CACHE_HIT, &OTLP_TRACE_AUX_CACHE_MISS), + ) + .await? + } else { + None + }; + // Requests without auxiliary admission first check the cache when preparing writes. + let aux_cache_metrics = aux_limiter + .is_none() + .then_some((&*OTLP_TRACE_AUX_CACHE_HIT, &*OTLP_TRACE_AUX_CACHE_MISS)); let is_trace_v1_model = matches!(pipeline, PipelineWay::OtlpTraceDirectV1); let data_model = match pipeline { PipelineWay::OtlpTraceDirectV1 => Some(TABLE_DATA_MODEL_TRACE_V1), @@ -918,7 +1116,11 @@ impl Instance { table_name: &table_name, }; let mut ingest_state = TraceIngestState { + // Confirmed hits stay satisfied for this request even if evicted + // later: cache eviction does not delete the stored lookup rows. + needs_aux_write: aux_limiter.is_none() || aux_permit.is_some(), aux_data: TraceAuxData::default(), + pending_aux_data: Vec::new(), outcome: TraceIngestOutcome::default(), failure_messages: TraceFailureMessages::default(), }; @@ -988,20 +1190,64 @@ impl Instance { OTLP_TRACES_ROWS.inc_by(ingest_state.outcome.accepted_spans as u64); + if !ingest_state.pending_aux_data.is_empty() { + let pending = std::mem::take(&mut ingest_state.pending_aux_data); + let instance = self.clone(); + let pipeline = match pipeline { + PipelineWay::OtlpTraceDirectV1 => PipelineWay::OtlpTraceDirectV1, + PipelineWay::OtlpTraceDirectV2 => PipelineWay::OtlpTraceDirectV2, + _ => PipelineWay::OtlpTraceDirectV0, + }; + let table_name = table_name.clone(); + let ctx = ctx.clone(); + common_runtime::spawn_ingest(async move { + let _permit = aux_permit; + let aux_data = confirmed_trace_aux_data(pending).await; + let result: ServerResult<()> = async { + let (requests, keys) = instance.prepare_trace_aux_write( + aux_data, + &pipeline, + &table_name, + &ctx, + aux_cache_metrics, + )?; + if !requests.inserts.is_empty() { + instance + .insert_trace_aux_requests(requests, keys, &pipeline, ctx) + .await?; + } + Ok(()) + } + .await; + if let Err(err) = result { + OTLP_TRACES_FAILURE_COUNT + .with_label_values(&["aux_table_update_failed"]) + .inc(); + warn!(table_name, err = ?err, "Failed to write auxiliary trace rows after asynchronous main write"); + } + }); + } + if !ingest_state.aux_data.is_empty() { - // Auxiliary trace tables are derived from spans whose main-table - // writes are already confirmed, so they never create new accepted - // spans and they do not affect rejected span counts. - let aux_requests = otlp::trace::to_grpc_insert_requests_for_aux_tables( + // Auxiliary trace tables are derived from accepted spans, so they + // never create new accepted spans or affect rejected span counts. + let aux_requests = self.prepare_trace_aux_write( std::mem::take(&mut ingest_state.aux_data), ingest_ctx.pipeline, ingest_ctx.table_name, + &ctx, + aux_cache_metrics, ); match aux_requests { - Ok((aux_requests, _)) if !aux_requests.inserts.is_empty() => { + Ok((aux_requests, pending_aux_keys)) if !aux_requests.inserts.is_empty() => { match self - .insert_trace_requests(aux_requests, ingest_ctx.pipeline, ctx) + .insert_trace_aux_requests( + aux_requests, + pending_aux_keys, + ingest_ctx.pipeline, + ctx, + ) .await { Ok(output) => { @@ -1076,6 +1322,72 @@ impl Instance { Ok(ingest_state.outcome) } + fn prepare_trace_aux_write( + &self, + mut aux_data: TraceAuxData, + pipeline: &PipelineWay, + table_name: &str, + ctx: &QueryContextRef, + metrics: Option<(&IntCounter, &IntCounter)>, + ) -> ServerResult<(RowInsertRequests, Vec)> { + let keys = self.trace_aux_cache.filter( + TableName::new(ctx.current_catalog(), ctx.current_schema(), table_name), + &mut aux_data, + metrics, + ); + if aux_data.is_empty() { + return Ok((RowInsertRequests::default(), keys)); + } + let (mut requests, _) = + otlp::trace::to_grpc_insert_requests_for_aux_tables(aux_data, pipeline, table_name)?; + // Remove fully cached tables before v1 reconciliation performs catalog lookups. + requests.inserts.retain(|request| { + request + .rows + .as_ref() + .is_some_and(|rows| !rows.rows.is_empty()) + }); + Ok((requests, keys)) + } + + /// Checks deferred cache population without confusing visible rows with completion. + #[cfg(feature = "testing")] + pub fn trace_aux_data_cached_for_test( + &self, + table_name: &str, + mut aux_data: TraceAuxData, + ctx: &QueryContextRef, + ) -> bool { + self.trace_aux_cache + .filter( + TableName::new(ctx.current_catalog(), ctx.current_schema(), table_name), + &mut aux_data, + None, + ) + .is_empty() + } + + async fn insert_trace_aux_requests( + &self, + requests: RowInsertRequests, + keys: Vec, + pipeline: &PipelineWay, + ctx: QueryContextRef, + ) -> ServerResult { + // Cache entries require a completed write, even when main spans use + // batching with early acknowledgement. + let mut aux_ctx = ctx.fork(); + aux_ctx.set_batching_enabled(false); + let output = self + .insert_trace_requests(requests, pipeline, Arc::new(aux_ctx)) + .await?; + // ponytail: cache skip_wal writes too. After datanode data loss, lookup + // rows may stay missing until eviction/restart; require durable writes + // before caching if that limitation becomes unacceptable. + self.trace_aux_cache.record(keys); + Ok(output) + } + /// Ingest one owned trace chunk so successful spans can be moved into the /// accepted set without extra cloning. async fn ingest_trace_chunk( @@ -1167,8 +1479,12 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += chunk_rows; - for span in &chunk { - ingest_state.aux_data.observe_span(span); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + for span in &chunk { + aux_data.observe_span(span); + } + ingest_state.add_aux_data(aux_data, output.meta.write_completions); } Ok(()) @@ -1345,7 +1661,11 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += span_count; - retry.add_to_aux_data(&mut ingest_state.aux_data); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + retry.add_to_aux_data(&mut aux_data); + ingest_state.add_aux_data(aux_data, output.meta.write_completions); + } Ok(()) } @@ -1428,7 +1748,11 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += 1; - span.add_to_aux_data(&mut ingest_state.aux_data); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + span.add_to_aux_data(&mut aux_data); + ingest_state.add_aux_data(aux_data, output.meta.write_completions); + } } Ok(()) diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index 1459119efdd..d497db83e13 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use api::v1::column_data_type_extension::TypeExt; use api::v1::value::ValueData; @@ -20,26 +21,269 @@ use api::v1::{ ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, Row, RowInsertRequest, Rows, SemanticType, Value, }; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::WriteCompletion; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ ColumnSchema as DatatypesColumnSchema, SchemaBuilder as DatatypesSchemaBuilder, }; use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; use opentelemetry_proto::tonic::common::v1::{AnyValue, ArrayValue, KeyValue}; -use servers::otlp::trace::SERVICE_NAME_COLUMN; +use prometheus::IntCounter; use servers::otlp::trace::attributes::Attributes; use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; +use servers::otlp::trace::{SERVICE_NAME_COLUMN, TraceAuxData}; +use table::table_name::TableName; use super::{ - ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, - TraceRequestSchema, TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, + ChunkFailureReaction, Instance, TraceAuxCache, TraceAuxCacheKey, TraceAuxEntry, + TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, TraceRequestSchema, + TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, wrap_trace_alter_failure, }; use crate::metrics::OTLP_TRACES_FAILURE_COUNT; +fn trace_aux_data(service: &str, span: &str, kind: &str) -> TraceAuxData { + TraceAuxData { + services: HashSet::from([service.to_string()]), + operations: HashSet::from([(service.to_string(), span.to_string(), kind.to_string())]), + } +} + +fn trace_aux_cache_metrics() -> (IntCounter, IntCounter) { + ( + IntCounter::new("test_trace_aux_cache_hit", "Cache hits").unwrap(), + IntCounter::new("test_trace_aux_cache_miss", "Cache misses").unwrap(), + ) +} + +#[tokio::test] +async fn test_trace_aux_waits_for_successful_chunk_completions() { + let (success_tx, success_rx) = tokio::sync::oneshot::channel(); + let (failure_tx, failure_rx) = tokio::sync::oneshot::channel(); + let completion = |receiver: tokio::sync::oneshot::Receiver| { + WriteCompletion::new(async move { + if receiver.await.unwrap_or(false) { + Ok(()) + } else { + Err(common_error::ext::BoxedError::new( + servers::error::InternalSnafu { + err_msg: "main write failed", + } + .build(), + )) + } + }) + }; + let mut confirmed = Box::pin(super::confirmed_trace_aux_data(vec![ + ( + trace_aux_data("good", "span", "server"), + vec![completion(success_rx)], + ), + ( + trace_aux_data("bad", "span", "server"), + vec![completion(failure_rx)], + ), + ])); + assert!(futures::poll!(confirmed.as_mut()).is_pending()); + success_tx.send(true).unwrap(); + assert!(futures::poll!(confirmed.as_mut()).is_pending()); + // A lost completion sender must not count as successful storage. + drop(failure_tx); + let aux = confirmed.await; + assert_eq!(aux.services, HashSet::from(["good".to_string()])); + assert_eq!( + aux.operations, + trace_aux_data("good", "span", "server").operations + ); +} + +#[test] +fn test_trace_aux_cache_confirmed_writes() { + let cache = TraceAuxCache::new(1024); + let (hits, misses) = trace_aux_cache_metrics(); + let metrics = Some((&hits, &misses)); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + let pending = cache.filter(table.clone(), &mut data, metrics); + assert_eq!(pending.len(), 2); + assert_eq!((hits.get(), misses.get()), (0, 2)); + + // Preparing a write must not suppress a retry after failure or cancellation. + let mut retry = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(table.clone(), &mut retry, None).len(), 2); + assert_eq!((hits.get(), misses.get()), (0, 2)); + cache.record(pending); + + let mut repeated = trace_aux_data("svc", "span", "server"); + assert!( + cache + .clone() + .filter(table.clone(), &mut repeated, metrics) + .is_empty() + ); + assert!(repeated.is_empty()); + assert_eq!((hits.get(), misses.get()), (2, 2)); + cache.filter(table.clone(), &mut retry, None); + assert_eq!((hits.get(), misses.get()), (2, 2)); + + for (span, kind) in [("new_span", "server"), ("span", "client")] { + let mut data = trace_aux_data("svc", span, kind); + assert_eq!(cache.filter(table.clone(), &mut data, metrics).len(), 1); + assert!(data.services.is_empty()); + assert_eq!(data.operations.len(), 1); + } + assert_eq!((hits.get(), misses.get()), (4, 4)); + cache.filter(table.clone(), &mut TraceAuxData::default(), metrics); + assert_eq!((hits.get(), misses.get()), (4, 4)); + + let mut other_service = trace_aux_data("other", "span", "server"); + assert_eq!(cache.filter(table, &mut other_service, None).len(), 2); + + for other_table in [ + TableName::new("other", "public", "traces"), + TableName::new("greptime", "other", "traces"), + TableName::new("greptime", "public", "other"), + ] { + let mut data = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(other_table, &mut data, None).len(), 2); + } + + // An operation hit must not suppress a missing service entry. + let cache = TraceAuxCache::new(1024); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + data.services.clear(); + cache.record(cache.filter(table.clone(), &mut data, None)); + + let mut data = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(table, &mut data, None).len(), 1); + assert_eq!(data.services.len(), 1); + assert!(data.operations.is_empty()); +} + +#[test] +fn test_trace_aux_cache_zero_size_disables_caching() { + let cache = TraceAuxCache::new(0); + let (hits, misses) = trace_aux_cache_metrics(); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + cache.record(cache.filter(table.clone(), &mut data, Some((&hits, &misses)))); + + let mut repeated = trace_aux_data("svc", "span", "server"); + assert_eq!( + cache + .filter(table, &mut repeated, Some((&hits, &misses))) + .len(), + 2 + ); + assert_eq!(repeated.services.len(), 1); + assert_eq!(repeated.operations.len(), 1); + assert_eq!(hits.get(), 0); + assert_eq!(misses.get(), 4); +} + +#[test] +fn test_trace_aux_cache_accounts_for_key_memory() { + let cache = TraceAuxCache::new(1024); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + cache.record(cache.filter(table.clone(), &mut data, None)); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 2); + assert!(cache.entries.weighted_size() > 2); + assert!(cache.entries.weighted_size() <= 1024); + + // Spare string capacity counts even when the string content is empty. + let large = || String::with_capacity(2048); + let service = || TraceAuxEntry::Service("svc".to_string()); + for (table, entry) in [ + (TableName::new(large(), "public", "traces"), service()), + (TableName::new("greptime", large(), "traces"), service()), + (TableName::new("greptime", "public", large()), service()), + (table.clone(), TraceAuxEntry::Service(large())), + ( + table.clone(), + TraceAuxEntry::Operation(large(), String::new(), String::new()), + ), + ( + table.clone(), + TraceAuxEntry::Operation(String::new(), large(), String::new()), + ), + ( + table, + TraceAuxEntry::Operation(String::new(), String::new(), large()), + ), + ] { + let cache = TraceAuxCache::new(1024); + cache.record(vec![TraceAuxCacheKey { + table: Arc::new(table), + entry, + }]); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 0); + } + + // Empty strings still retain the key and table structs. + let cache = TraceAuxCache::new(1); + cache.record(vec![TraceAuxCacheKey { + table: Arc::new(TableName::new("", "", "")), + entry: TraceAuxEntry::Service(String::new()), + }]); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 0); +} + +#[tokio::test] +async fn test_trace_aux_admission_rechecks_cache() { + let cache = TraceAuxCache::new(1024); + let (hits, misses) = trace_aux_cache_metrics(); + let limiter = RequestLimiter::try_new(1).unwrap(); + let table = TableName::new("greptime", "public", "traces"); + let data = || trace_aux_data("svc", "span", "server"); + let permit = cache + .acquire_for_misses(&limiter, table.clone(), data(), (&hits, &misses)) + .await + .unwrap() + .unwrap(); + let cloned_cache = cache.clone(); + let cloned_limiter = limiter.clone(); + let mut waiting = Box::pin(cloned_cache.acquire_for_misses( + &cloned_limiter, + table.clone(), + data(), + (&hits, &misses), + )); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + assert_eq!((hits.get(), misses.get()), (0, 4)); + + // Successful storage populates the shared cache before releasing admission. + cache.record(cache.filter(table.clone(), &mut data(), None)); + assert!( + cloned_cache + .acquire_for_misses(&cloned_limiter, table.clone(), data(), (&hits, &misses)) + .await + .unwrap() + .is_none() + ); + drop(permit); + assert!(waiting.await.unwrap().is_none()); + assert_eq!((hits.get(), misses.get()), (2, 4)); + + // Rechecking a now-warm request must release its briefly acquired slot. + let other = cache.acquire_for_misses( + &limiter, + table, + trace_aux_data("svc", "other", "server"), + (&hits, &misses), + ); + assert!(futures::poll!(Box::pin(other).as_mut()).is_ready()); + assert_eq!((hits.get(), misses.get()), (3, 5)); +} + #[test] fn test_chunk_owned() { let chunks = chunk_owned(vec![1, 2, 3], 2); diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index 0295d12e685..6f37f14e34a 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -519,31 +519,34 @@ mod tests { ); } - /// The PHP instrumentation case: an empty string must be distinguishable - /// from any other unparsable value in the reported diagnostics. + /// The PHP instrumentation case: `getHeaderLine('Content-Length')` sends "" + /// when the header is absent. The span must be kept with a null attribute. #[test] - fn test_prepare_trace_column_rewrites_reports_empty_string_value() { - let rows = vec![Row { - values: vec![Value { - value_data: Some(ValueData::StringValue(String::new())), + fn test_prepare_trace_column_rewrites_nulls_empty_string_value() { + let mut rows = Rows { + schema: vec![ColumnSchema { + column_name: "span_attributes.http.response.body.size".to_string(), + datatype: ColumnDataType::String as i32, + ..Default::default() }], - }]; + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::StringValue(String::new())), + }], + }], + }; let pending_rewrites = vec![PendingTraceColumnRewrite { col_idx: 0, target_type: ColumnDataType::Int64, column_name: "span_attributes.http.response.body.size".to_string(), }]; - let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "opentelemetry_traces") - .unwrap_err(); - assert!( - err.error.to_string().contains( - "'span_attributes.http.response.body.size' in table 'opentelemetry_traces' \ - from String(\"\") to Int64" - ), - "unexpected error message: {}", - err.error - ); + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "opentelemetry_traces") + .unwrap() + .apply(&mut rows); + + assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32); + assert_eq!(rows.rows[0].values[0].value_data, None); } #[test] diff --git a/src/frontend/src/lib.rs b/src/frontend/src/lib.rs index 1468fe7c0b2..7c68e445242 100644 --- a/src/frontend/src/lib.rs +++ b/src/frontend/src/lib.rs @@ -23,3 +23,5 @@ pub(crate) mod metrics; pub mod server; pub mod service_config; mod stream_wrapper; + +pub use metrics::{OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index f649db0b46f..74574e56930 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use lazy_static::lazy_static; use prometheus::*; @@ -68,6 +69,12 @@ lazy_static! { ) .unwrap(); + /// Initial trace auxiliary-cache lookups, excluding admission and write rechecks. + pub static ref OTLP_TRACE_AUX_CACHE_HIT: IntCounter = + CACHE_HIT.with_label_values(&["otlp_trace_aux"]); + pub static ref OTLP_TRACE_AUX_CACHE_MISS: IntCounter = + CACHE_MISS.with_label_values(&["otlp_trace_aux"]); + /// The number of OpenTelemetry logs send by frontend node. pub static ref OTLP_LOGS_ROWS: IntCounter = register_int_counter!( "greptime_frontend_otlp_logs_rows", diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 017bba53fba..8de7bfe340f 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -24,9 +24,7 @@ use common_base::Plugins; use common_config::Configurable; use common_telemetry::{info, warn}; use meta_client::MetaClientOptions; -use servers::batcher::logical_table::{ - LogicalTablePendingRowsBatcher, pending_rows_batch_sync_enabled, -}; +use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled}; use servers::error::Error as ServerError; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::FlightCraftRef; @@ -36,7 +34,7 @@ use servers::grpc::{GrpcOptions, GrpcServer}; use servers::http::event::LogValidatorRef; use servers::http::result::error_result::ErrorResponse; use servers::http::utils::router::RouterConfigurator; -use servers::http::{BatchingProtocol, HttpOptions, HttpServer, HttpServerBuilder}; +use servers::http::{HttpOptions, HttpServer, HttpServerBuilder}; use servers::interceptor::LogIngestInterceptorRef; use servers::metrics_handler::MetricsHandler; use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef}; @@ -74,6 +72,7 @@ where { pub fn new(opts: T, instance: Arc, plugins: Plugins) -> Self { let feopts = opts.clone().into(); + instance.init_logical_batcher(&feopts); // Create server request memory limiter for all server protocols let server_memory_limiter = ServerMemoryLimiter::new( feopts.max_in_flight_write_bytes.as_bytes(), @@ -110,7 +109,8 @@ where request_memory_limiter: ServerMemoryLimiter, ) -> HttpServerBuilder { let mut builder = HttpServerBuilder::new(effective_http_options(opts)) - .with_batching_protocols(opts.pending_rows_batcher.protocols.clone()) + .with_batching_protocols(opts.table_batcher_options().protocols.clone()) + .with_logical_batching_protocols(opts.logical_batcher_options().protocols) .with_memory_limiter(request_memory_limiter) .with_sql_handler(self.instance.clone()); @@ -134,43 +134,26 @@ where let prom_store = effective_prom_store_options(opts); if prom_store.enable { - let pending_rows_batcher = if prom_store.with_metric_engine { - LogicalTablePendingRowsBatcher::try_new( - self.instance.partition_manager().clone(), - self.instance.node_manager().clone(), - self.instance.catalog_manager().clone(), - self.instance.table_flownode_set_cache().clone(), - prom_store.with_metric_engine, - self.instance.clone(), - prom_store.pending_rows_flush_interval, - prom_store.max_batch_rows, - prom_store.max_concurrent_flushes, - prom_store.worker_channel_capacity, - prom_store.max_inflight_requests, - prom_store.flow_notification_queue_capacity, - ) - } else { - None - }; + let pending_rows_batcher = opts + .logical_batcher_options() + .protocols + .contains(&BatchingProtocol::Prom) + .then(|| self.instance.logical_batcher().cloned()) + .flatten(); builder = builder .with_prom_handler( self.instance.clone(), Some(self.instance.clone()), opts.prom_store.with_metric_engine, opts.prom_store.prom_validation_mode, - opts.prom_store - .experimental_enable_prometheus_native_histogram, pending_rows_batcher, ) .with_prometheus_handler(self.instance.clone()); } if opts.otlp.enable { - builder = builder.with_otlp_handler( - self.instance.clone(), - opts.prom_store.with_metric_engine, - opts.otlp.experimental_enable_exponential_histogram, - ); + builder = builder + .with_otlp_handler(self.instance.clone(), opts.prom_store.with_metric_engine); } if opts.jaeger.enable { @@ -232,11 +215,15 @@ where external: bool, request_memory_limiter: ServerMemoryLimiter, ) -> Result { - let builder = if let Some(builder) = self.grpc_server_builder.take() { + let mut builder = if let Some(builder) = self.grpc_server_builder.take() { builder } else { self.grpc_server_builder(grpc, request_memory_limiter)? }; + // Browsers never talk to the internal server. + if external && grpc.enable_cors { + builder = builder.with_cors(grpc.cors_allowed_origins.clone()); + } let user_provider = if external { self.plugins.get::() @@ -381,6 +368,15 @@ where } } + let table_batcher = opts.table_batcher_options(); + let batching_enabled = table_batcher.pending_rows_batching_enabled(); + let mysql_batching = + batching_enabled && table_batcher.protocols.contains(&BatchingProtocol::Mysql); + let postgres_batching = batching_enabled + && table_batcher + .protocols + .contains(&BatchingProtocol::Postgres); + if opts.mysql.enable { // Init MySQL server let opts = &opts.mysql; @@ -396,13 +392,16 @@ where let mysql_server = MysqlServer::create_server( common_runtime::global_runtime(), Arc::new(MysqlSpawnRef::new(instance.clone(), user_provider.clone())), - Arc::new(MysqlSpawnConfig::new( - opts.tls.should_force_tls(), - tls_server_config, - opts.keep_alive.as_secs(), - opts.reject_no_database.unwrap_or(false), - opts.prepared_stmt_cache_size, - )), + Arc::new( + MysqlSpawnConfig::new( + opts.tls.should_force_tls(), + tls_server_config, + opts.keep_alive.as_secs(), + opts.reject_no_database.unwrap_or(false), + opts.prepared_stmt_cache_size, + ) + .with_batching_enabled(mysql_batching), + ), Some(instance.process_manager().clone()), ); handlers.insert((mysql_server, mysql_addr)); @@ -419,15 +418,18 @@ where maybe_watch_server_tls_config(tls_server_config.clone()).context(StartServerSnafu)?; - let pg_server = Box::new(PostgresServer::new( - instance.clone(), - opts.tls.should_force_tls(), - tls_server_config, - opts.keep_alive.as_secs(), - common_runtime::global_runtime(), - user_provider.clone(), - Some(self.instance.process_manager().clone()), - )) as Box; + let pg_server = Box::new( + PostgresServer::new( + instance.clone(), + opts.tls.should_force_tls(), + tls_server_config, + opts.keep_alive.as_secs(), + common_runtime::global_runtime(), + user_provider.clone(), + Some(self.instance.process_manager().clone()), + ) + .with_batching_enabled(postgres_batching), + ) as Box; handlers.insert((pg_server, pg_addr)); } @@ -439,38 +441,54 @@ where /// Selected shared controls override legacy Prom batching knobs, not protocol behavior. fn effective_prom_store_options(opts: &FrontendOptions) -> PromStoreOptions { let mut prom_store = opts.prom_store.clone(); - let shared = &opts.pending_rows_batcher; - if shared.protocols.contains(&BatchingProtocol::Prom) && shared.pending_rows_batching_enabled() - { + let shared = opts.logical_batcher_options(); + if shared.protocols.contains(&BatchingProtocol::Prom) { prom_store.pending_rows_flush_interval = shared.pending_rows_flush_interval; prom_store.max_batch_rows = shared.max_batch_rows; prom_store.max_concurrent_flushes = shared.max_concurrent_flushes; prom_store.worker_channel_capacity = shared.worker_channel_capacity; prom_store.max_inflight_requests = shared.max_inflight_requests; prom_store.flow_notification_queue_capacity = shared.flow_notification_queue_capacity; + } else { + prom_store.pending_rows_flush_interval = Duration::ZERO; } prom_store } -fn effective_http_options(opts: &FrontendOptions) -> HttpOptions { +/// Resolves HTTP options, allowing synchronous batches enough time to flush. +pub fn effective_http_options(opts: &FrontendOptions) -> HttpOptions { effective_http_options_with_sync(opts, pending_rows_batch_sync_enabled()) } fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> HttpOptions { let mut http = opts.http.clone(); let prom_store = effective_prom_store_options(opts); - let shared = &opts.pending_rows_batcher; - // Ordinary-table batching always waits for its flush, independently of the - // dedicated Prom batcher's asynchronous acknowledgement mode. - let common_enabled = shared.pending_rows_batching_enabled() + let shared = opts.table_batcher_options(); + let common_enabled = batch_sync + && shared.pending_rows_batching_enabled() && shared.protocols.iter().any(|protocol| { - *protocol != BatchingProtocol::Prom - || (prom_store.enable && !prom_store.with_metric_engine) + !matches!( + protocol, + BatchingProtocol::Mysql | BatchingProtocol::Postgres + ) && (*protocol != BatchingProtocol::Prom + || (prom_store.enable && !prom_store.with_metric_engine)) }); let common_interval = common_enabled.then_some(shared.pending_rows_flush_interval); let prom_interval = (prom_store.pending_rows_batching_enabled() && batch_sync) .then_some(prom_store.pending_rows_flush_interval); - let Some(flush_interval) = common_interval.into_iter().chain(prom_interval).max() else { + let logical = opts.logical_batcher_options(); + let otlp_interval = (batch_sync + && opts.otlp.enable + && opts.prom_store.with_metric_engine + && logical.protocols.contains(&BatchingProtocol::Otlp) + && logical.pending_rows_batching_enabled()) + .then_some(logical.pending_rows_flush_interval); + let Some(flush_interval) = common_interval + .into_iter() + .chain(prom_interval) + .chain(otlp_interval) + .max() + else { return http; }; let fallback_timeout = flush_interval.saturating_add(Duration::from_secs(1)); @@ -507,7 +525,9 @@ mod tests { use async_trait::async_trait; use auth::{UserProviderRef, static_user_provider_from_option}; use client::{Client, Database}; + use common_grpc::channel_manager::ChannelManager; use meta_client::client::MetaClientBuilder; + use reqwest::header::{ACCESS_CONTROL_ALLOW_ORIGIN, ACCESS_CONTROL_REQUEST_METHOD, ORIGIN}; use servers::grpc::GRPC_SERVER; use servers::grpc::flight::{FlightCraft, FlightCraftRef, TonicStream}; use tonic::{Code, Request, Response, Status, Streaming}; @@ -515,6 +535,86 @@ mod tests { use crate::instance::builder::FrontendBuilder; use crate::server::*; + #[tokio::test] + async fn test_logical_batcher_shared_without_prom_endpoint() { + use crate::service_config::BatcherOptions; + let mut options = FrontendOptions::default(); + options.prom_store.enable = false; + options.pending_rows_batcher.logical_table = Some(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }); + let meta_client = Arc::new( + MetaClientBuilder::new(0, Role::Frontend) + .enable_procedure() + .build(), + ); + let instance = Arc::new( + FrontendBuilder::new_test(&options, meta_client) + .try_build() + .await + .unwrap(), + ); + let mut disabled_options = options.clone(); + disabled_options.otlp.enable = false; + instance.init_logical_batcher(&disabled_options); + assert!(instance.logical_batcher().is_none()); + + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let batcher = instance.logical_batcher().unwrap().clone(); + instance.init_logical_batcher(&options); + assert!(Arc::ptr_eq(&batcher, instance.logical_batcher().unwrap())); + let weak = Arc::downgrade(&instance); + drop(services); + drop(instance); + assert!( + weak.upgrade().is_none(), + "batcher must not retain its schema owner" + ); + } + + #[test] + fn test_logical_batcher_http_timeout_and_prom_disable() { + use crate::service_config::pending_rows_batcher::BatcherOptions; + let mut opts = FrontendOptions::default(); + opts.http.timeout = Duration::from_secs(1); + opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2); + opts.pending_rows_batcher.logical_table = Some(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(5), + ..Default::default() + }); + assert!(!effective_prom_store_options(&opts).pending_rows_batching_enabled()); + // Both logical protocols follow the global acknowledgement policy, + // independently of enabling the Prom HTTP endpoint. + opts.prom_store.enable = false; + assert_eq!( + effective_http_options_with_sync(&opts, true).timeout, + Duration::from_secs(6) + ); + assert_eq!( + effective_http_options_with_sync(&opts, false).timeout, + opts.http.timeout + ); + opts.prom_store.with_metric_engine = false; + assert_eq!( + effective_http_options_with_sync(&opts, false).timeout, + opts.http.timeout + ); + opts.prom_store.with_metric_engine = true; + opts.pending_rows_batcher + .logical_table + .as_mut() + .unwrap() + .protocols + .clear(); + assert_eq!( + effective_http_options_with_sync(&opts, true).timeout, + opts.http.timeout + ); + } + #[test] fn test_effective_prom_batching_controls() { // Only an enabled shared Prom selection replaces the legacy controls. @@ -530,9 +630,7 @@ mod tests { opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2); opts.prom_store.with_metric_engine = metric_engine; opts.prom_store.enable = prom_enabled; - opts.prom_store - .experimental_enable_prometheus_native_histogram = true; - let shared = &mut opts.pending_rows_batcher; + let shared = &mut opts.pending_rows_batcher.table; shared.protocols = vec![if selected { BatchingProtocol::Prom } else { @@ -569,7 +667,7 @@ mod tests { #[test] fn test_http_timeout_covers_synchronous_batchers() { - // Shared ordinary writes remain synchronous even when Prom is asynchronous. + // Only synchronous batchers extend the HTTP timeout. for ( protocols, metric_engine, @@ -579,10 +677,14 @@ mod tests { timeout_secs, expected_secs, ) in [ - (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 6), + (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 1), + (vec![BatchingProtocol::Mysql], false, false, 5, 0, 1, 1), + (vec![BatchingProtocol::Postgres], false, false, 5, 0, 1, 1), + (vec![BatchingProtocol::Mysql], false, true, 5, 0, 1, 1), + (vec![BatchingProtocol::Postgres], false, true, 5, 0, 1, 1), (vec![BatchingProtocol::Prom], true, false, 5, 2, 1, 1), (vec![BatchingProtocol::Prom], true, true, 5, 2, 1, 6), - (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 6), + (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 1), (vec![BatchingProtocol::Influxdb], true, true, 5, 8, 1, 9), (vec![BatchingProtocol::Influxdb], true, true, 8, 5, 1, 9), (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 0, 0), @@ -595,8 +697,8 @@ mod tests { opts.http.timeout = Duration::from_secs(timeout_secs); opts.prom_store.with_metric_engine = metric_engine; opts.prom_store.pending_rows_flush_interval = Duration::from_secs(legacy_secs); - opts.pending_rows_batcher.protocols = protocols; - opts.pending_rows_batcher.pending_rows_flush_interval = + opts.pending_rows_batcher.table.protocols = protocols; + opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::from_secs(shared_secs); assert_eq!( effective_http_options_with_sync(&opts, batch_sync).timeout, @@ -697,24 +799,30 @@ mod tests { fn test_invalid_shared_batching_preserves_prom_store_options() { type KnobMutator = fn(&mut FrontendOptions); let cases: [KnobMutator; 5] = [ - |opts| opts.pending_rows_batcher.max_concurrent_flushes = usize::MAX, - |opts| opts.pending_rows_batcher.worker_channel_capacity = usize::MAX, - |opts| opts.pending_rows_batcher.max_inflight_requests = usize::MAX, + |opts| opts.pending_rows_batcher.table.max_concurrent_flushes = usize::MAX, + |opts| opts.pending_rows_batcher.table.worker_channel_capacity = usize::MAX, + |opts| opts.pending_rows_batcher.table.max_inflight_requests = usize::MAX, |opts| { - opts.pending_rows_batcher.flow_notification_queue_capacity = - NonZeroUsize::new(usize::MAX).unwrap() + opts.pending_rows_batcher + .table + .flow_notification_queue_capacity = NonZeroUsize::new(usize::MAX).unwrap() }, - |opts| opts.pending_rows_batcher.pending_rows_flush_interval = Duration::MAX, + |opts| opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::MAX, ]; for invalidate in cases { let mut opts = FrontendOptions::default(); opts.http.timeout = Duration::from_secs(1); opts.prom_store.pending_rows_flush_interval = Duration::from_secs(5); - opts.pending_rows_batcher.protocols = + opts.pending_rows_batcher.table.protocols = vec![BatchingProtocol::Prom, BatchingProtocol::Influxdb]; - opts.pending_rows_batcher.pending_rows_flush_interval = Duration::from_secs(10); + opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::from_secs(10); invalidate(&mut opts); - assert!(!opts.pending_rows_batcher.pending_rows_batching_enabled()); + assert!( + !opts + .pending_rows_batcher + .table + .pending_rows_batching_enabled() + ); assert_eq!(opts.prom_store, effective_prom_store_options(&opts)); assert_eq!( Duration::from_secs(6), @@ -842,12 +950,20 @@ mod tests { let public_database = Database::new( "greptime", "public", - Client::with_urls([public_addr.to_string()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [public_addr.to_string()], + ), ); let internal_database = Database::new( "greptime", "public", - Client::with_urls([internal_addr.to_string()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [internal_addr.to_string()], + ), ); let internal_result = internal_database.sql("SELECT 1").await; @@ -924,4 +1040,68 @@ mod tests { // Assert assert!(health_check.is_ok()); } + + #[tokio::test] + async fn test_internal_grpc_server_never_serves_cors() { + let options = FrontendOptions { + http: HttpOptions { + addr: "127.0.0.1:0".to_string(), + ..Default::default() + }, + grpc: GrpcOptions { + enable_cors: true, + ..GrpcOptions::default().with_bind_addr("127.0.0.1:0") + }, + internal_grpc: Some(GrpcOptions { + enable_cors: true, + ..GrpcOptions::default().with_bind_addr("127.0.0.1:0") + }), + mysql: crate::service_config::MysqlOptions { + enable: false, + ..Default::default() + }, + postgres: crate::service_config::PostgresOptions { + enable: false, + ..Default::default() + }, + ..Default::default() + }; + let meta_client = Arc::new( + MetaClientBuilder::new(0, Role::Frontend) + .enable_procedure() + .build(), + ); + let instance = Arc::new( + FrontendBuilder::new_test(&options, meta_client) + .try_build() + .await + .unwrap(), + ); + let mut services = Services::new(options, instance, Default::default()) + .build() + .unwrap(); + + services.start_all().await.unwrap(); + let public_addr = services.addr(GRPC_SERVER).unwrap(); + let internal_addr = services.addr("INTERNAL_GRPC_SERVER").unwrap(); + let public = send_cors_preflight(public_addr).await; + let internal = send_cors_preflight(internal_addr).await; + services.shutdown_all().await.unwrap(); + + assert!(public.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN)); + assert!(!internal.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN)); + } + + async fn send_cors_preflight(addr: std::net::SocketAddr) -> reqwest::Response { + reqwest::Client::new() + .request( + reqwest::Method::OPTIONS, + format!("http://{addr}/greptime.v1.HealthCheck/Check"), + ) + .header(ORIGIN, "https://example.com") + .header(ACCESS_CONTROL_REQUEST_METHOD, "POST") + .send() + .await + .unwrap() + } } diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs index bcd1097cfcf..2c761c4242a 100644 --- a/src/frontend/src/service_config.rs +++ b/src/frontend/src/service_config.rs @@ -26,6 +26,6 @@ pub use jaeger::JaegerOptions; pub use mysql::MysqlOptions; pub use opentsdb::OpentsdbOptions; pub use otlp::OtlpOptions; -pub use pending_rows_batcher::PendingRowsBatcherOptions; +pub use pending_rows_batcher::{BatcherOptions, PendingRowsBatcherOptions}; pub use postgres::PostgresOptions; pub use prom_store::PromStoreOptions; diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index d6a36f702e4..7a6d3f75b91 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -12,17 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_base::readable_size::ReadableSize; +use common_stat::get_total_memory_readable; use serde::{Deserialize, Serialize}; const DEFAULT_TRACE_INGEST_CHUNK_SIZE: usize = 512; +const MIN_TRACE_AUX_CACHE_SIZE: ReadableSize = ReadableSize::mb(32); #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] pub struct OtlpOptions { pub enable: bool, - pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, + /// Estimated memory budget for cached trace service/operation keys per frontend, + /// shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. + /// Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. + /// Uses 32MiB if memory detection is unavailable. + /// Excludes cache and allocator overhead; shared table names are charged per entry. + pub trace_aux_cache_size: ReadableSize, /// Whether to synthesize the `greptime_otel_resource_info` descriptor /// table from the resource attributes of OTLP metrics, so metrics-only /// services reach the semantic graph. Off by default: it creates and @@ -34,28 +42,52 @@ impl Default for OtlpOptions { fn default() -> Self { Self { enable: true, - experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, + trace_aux_cache_size: default_trace_aux_cache_size(get_total_memory_readable()), experimental_enable_resource_info: false, } } } +fn default_trace_aux_cache_size(memory: Option) -> ReadableSize { + memory + .map(|size| size / 128) + .unwrap_or_default() + .max(MIN_TRACE_AUX_CACHE_SIZE) +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn test_default_trace_aux_cache_size() { + for (memory, expected) in [ + (None, ReadableSize::mb(32)), + (Some(ReadableSize(0)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(1)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(4)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(8)), ReadableSize::mb(64)), + (Some(ReadableSize::gb(64)), ReadableSize::mb(512)), + ] { + assert_eq!(default_trace_aux_cache_size(memory), expected); + } + } + #[test] fn test_otlp_options() { let default = OtlpOptions::default(); assert!(default.enable); - assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); + assert_eq!( + default.trace_aux_cache_size, + default_trace_aux_cache_size(get_total_memory_readable()) + ); assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); - assert!(!options.experimental_enable_exponential_histogram); + assert_eq!(options.trace_aux_cache_size, default.trace_aux_cache_size); assert_eq!( options.trace_ingest_chunk_size, DEFAULT_TRACE_INGEST_CHUNK_SIZE @@ -63,15 +95,17 @@ mod tests { let options: OtlpOptions = toml::from_str("trace_ingest_chunk_size = 0").unwrap(); assert!(options.enable); - assert!(!options.experimental_enable_exponential_histogram); assert_eq!(options.trace_ingest_chunk_size, 0); - let options: OtlpOptions = - toml::from_str("experimental_enable_exponential_histogram = true").unwrap(); - assert!(options.experimental_enable_exponential_histogram); - let serialized = toml::to_string(&options).unwrap(); - assert!(serialized.contains("experimental_enable_exponential_histogram = true")); assert_eq!(toml::from_str::(&serialized).unwrap(), options); + + for (value, size) in [("0", ReadableSize(0)), ("\"1MiB\"", ReadableSize::mb(1))] { + let options: OtlpOptions = + toml::from_str(&format!("trace_aux_cache_size = {value}")).unwrap(); + assert_eq!(options.trace_aux_cache_size, size); + let serialized = toml::to_string(&options).unwrap(); + assert_eq!(toml::from_str::(&serialized).unwrap(), options); + } } } diff --git a/src/frontend/src/service_config/pending_rows_batcher.rs b/src/frontend/src/service_config/pending_rows_batcher.rs index ef7ce14ef02..b671ee9a272 100644 --- a/src/frontend/src/service_config/pending_rows_batcher.rs +++ b/src/frontend/src/service_config/pending_rows_batcher.rs @@ -17,14 +17,32 @@ use std::time::Duration; use common_batcher::flush_policy::timing::TimingFlushPolicy; use serde::{Deserialize, Serialize}; -use servers::http::BatchingProtocol; +use servers::batcher::BatchingProtocol; use tokio::sync::Semaphore; -/// Experimental table write batching options shared by HTTP ingestion protocols. -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +use crate::frontend::FrontendOptions; + +/// Independent ordinary-table and logical-table batching configuration. +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] pub struct PendingRowsBatcherOptions { - /// HTTP write protocols sharing this batcher; empty disables all entrances. + /// Existing ordinary-table controls retain their original TOML paths. + #[serde(flatten)] + pub table: BatcherOptions, + /// Independent logical-table controls; absence retains the legacy Prom fallback. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_logical_options" + )] + pub logical_table: Option, +} + +/// Write batching controls shared by ingestion protocols. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default)] +pub struct BatcherOptions { + /// Write protocols sharing this batcher; empty disables all entrances. pub protocols: Vec, /// Time from the first pending submission to a timed flush. Zero disables batching. #[serde(with = "humantime_serde")] @@ -41,7 +59,7 @@ pub struct PendingRowsBatcherOptions { pub flow_notification_queue_capacity: NonZeroUsize, } -impl PendingRowsBatcherOptions { +impl BatcherOptions { /// Returns whether a protocol opts in and its controls pass construction validation. pub fn pending_rows_batching_enabled(&self) -> bool { !self.protocols.is_empty() @@ -58,79 +76,263 @@ impl PendingRowsBatcherOptions { } } -impl Default for PendingRowsBatcherOptions { +impl Default for BatcherOptions { fn default() -> Self { Self { protocols: Vec::new(), pending_rows_flush_interval: Duration::ZERO, max_batch_rows: 100_000, max_concurrent_flushes: 256, - worker_channel_capacity: 65526, + worker_channel_capacity: 65_536, max_inflight_requests: 3000, flow_notification_queue_capacity: NonZeroUsize::new(1024).unwrap_or(NonZeroUsize::MIN), } } } +/// Rejects protocols that cannot write metric-engine logical tables. +fn deserialize_logical_options<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let options = Option::::deserialize(deserializer)?; + if options.as_ref().is_some_and(|options| { + options + .protocols + .iter() + .any(|protocol| !matches!(protocol, BatchingProtocol::Prom | BatchingProtocol::Otlp)) + }) { + return Err(serde::de::Error::custom( + "pending_rows_batcher.logical_table only supports prom and otlp", + )); + } + Ok(options) +} + +impl FrontendOptions { + /// Explicit configuration wins even when it disables batching. + pub(crate) fn table_batcher_options(&self) -> &BatcherOptions { + &self.pending_rows_batcher.table + } + + /// Resolve legacy Prom controls once without merging fields across config blocks. + pub(crate) fn logical_batcher_options(&self) -> BatcherOptions { + if let Some(options) = &self.pending_rows_batcher.logical_table { + return options.clone(); + } + let shared = &self.pending_rows_batcher.table; + if shared.protocols.contains(&BatchingProtocol::Prom) + && shared.pending_rows_batching_enabled() + { + let mut options = shared.clone(); + options.protocols = vec![BatchingProtocol::Prom]; + return options; + } + let prom = &self.prom_store; + BatcherOptions { + protocols: vec![BatchingProtocol::Prom], + pending_rows_flush_interval: prom.pending_rows_flush_interval, + max_batch_rows: prom.max_batch_rows, + max_concurrent_flushes: prom.max_concurrent_flushes, + worker_channel_capacity: prom.worker_channel_capacity, + max_inflight_requests: prom.max_inflight_requests, + flow_notification_queue_capacity: prom.flow_notification_queue_capacity, + } + } +} + #[cfg(test)] mod tests { use crate::service_config::pending_rows_batcher::*; + #[test] + fn test_batcher_config_precedence() { + use crate::frontend::FrontendOptions; + let legacy = "[prom_store]\nenable = true\nwith_metric_engine = true\npending_rows_flush_interval = '2s'\nmax_batch_rows = 9\n"; + let shared = "[pending_rows_batcher]\nprotocols = ['prom', 'otlp']\npending_rows_flush_interval = '3s'\n"; + let options: FrontendOptions = toml::from_str(legacy).unwrap(); + assert_eq!(options.logical_batcher_options().max_batch_rows, 9); + assert_eq!( + options.logical_batcher_options().protocols, + vec![BatchingProtocol::Prom] + ); + let options: FrontendOptions = toml::from_str(&format!("{legacy}{shared}")).unwrap(); + assert_eq!( + options + .logical_batcher_options() + .pending_rows_flush_interval, + Duration::from_secs(3) + ); + assert_eq!(options.logical_batcher_options().max_batch_rows, 100_000); + assert_eq!( + options.logical_batcher_options().protocols, + vec![BatchingProtocol::Prom] + ); + for explicit in [ + "", + "protocols = ['otlp']\npending_rows_flush_interval = '5s'", + ] { + let options: FrontendOptions = toml::from_str(&format!( + "{legacy}{shared}[pending_rows_batcher.logical_table]\n{explicit}" + )) + .unwrap(); + assert_eq!( + &options.logical_batcher_options(), + options.pending_rows_batcher.logical_table.as_ref().unwrap() + ); + assert_eq!( + options.table_batcher_options(), + &toml::from_str::(&format!("{legacy}{shared}")) + .unwrap() + .pending_rows_batcher + .table + ); + let restored: FrontendOptions = + toml::from_str(&toml::to_string(&options).unwrap()).unwrap(); + assert_eq!( + options.pending_rows_batcher.logical_table, + restored.pending_rows_batcher.logical_table + ); + assert_eq!(options.pending_rows_batcher, restored.pending_rows_batcher); + } + } + + #[test] + fn test_batcher_config_independent_defaults() { + use crate::frontend::FrontendOptions; + + let parent = "[pending_rows_batcher]\nprotocols = ['prom']\npending_rows_flush_interval = '3s'\nmax_batch_rows = 7\n"; + let options: FrontendOptions = toml::from_str(parent).unwrap(); + let serialized = toml::to_string(&options).unwrap(); + assert!(!serialized.contains("[pending_rows_batcher.logical_table]")); + assert!(!serialized.contains("[pending_rows_batcher.table]")); + let restored: FrontendOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(restored.logical_batcher_options().max_batch_rows, 7); + + let options: FrontendOptions = + toml::from_str(&format!("{parent}[pending_rows_batcher.logical_table]")).unwrap(); + assert_eq!(options.table_batcher_options().max_batch_rows, 7); + assert_eq!(options.logical_batcher_options(), BatcherOptions::default()); + assert!( + !options + .logical_batcher_options() + .pending_rows_batching_enabled() + ); + } + + #[test] + fn test_logical_batcher_config_protocols() { + use crate::frontend::FrontendOptions; + for protocols in ["[]", "['prom']", "['otlp']", "['prom', 'otlp']"] { + let options: FrontendOptions = toml::from_str(&format!( + "[pending_rows_batcher.logical_table]\nprotocols = {protocols}" + )) + .unwrap(); + assert!(options.pending_rows_batcher.logical_table.is_some()); + } + for protocol in [ + "influxdb", + "logs", + "loki", + "http_sql", + "opentsdb", + "elasticsearch", + "splunk", + "mysql", + "postgres", + ] { + assert!( + toml::from_str::(&format!( + "[pending_rows_batcher.logical_table]\nprotocols = ['{protocol}']" + )) + .is_err() + ); + } + } + #[test] fn test_protocols() { - let options: PendingRowsBatcherOptions = toml::from_str( - "protocols = ['influxdb', 'opentsdb', 'otlp', 'logs', 'loki', 'splunk', 'elasticsearch', 'http_sql', 'prom']", + let options: BatcherOptions = toml::from_str( + "protocols = ['influxdb', 'opentsdb', 'otlp', 'logs', 'loki', 'splunk', 'elasticsearch', 'http_sql', 'prom', 'mysql', 'postgres']", ).unwrap(); - assert_eq!(options.protocols.len(), 9); + assert_eq!(options.protocols.len(), 11); + assert!(options.protocols.contains(&BatchingProtocol::Mysql)); + assert!(options.protocols.contains(&BatchingProtocol::Postgres)); assert!(options.protocols.contains(&BatchingProtocol::HttpSql)); - assert!(PendingRowsBatcherOptions::default().protocols.is_empty()); + assert!(BatcherOptions::default().protocols.is_empty()); for invalid in ["sql", "jaeger", "unknown"] { assert!( - toml::from_str::(&format!("protocols = ['{invalid}']")) - .is_err() + toml::from_str::(&format!("protocols = ['{invalid}']")).is_err() ); } } #[test] fn test_notification_capacity() { - let default = PendingRowsBatcherOptions::default(); + let default = BatcherOptions::default(); assert_eq!(default.flow_notification_queue_capacity.get(), 1024); - let configured: PendingRowsBatcherOptions = + let configured: BatcherOptions = toml::from_str("flow_notification_queue_capacity = 8").unwrap(); assert_eq!(configured.flow_notification_queue_capacity.get(), 8); - assert!( - toml::from_str::("flow_notification_queue_capacity = 0") - .is_err() - ); + assert!(toml::from_str::("flow_notification_queue_capacity = 0").is_err()); } #[test] fn test_defaults_and_roundtrip() { - let options: PendingRowsBatcherOptions = toml::from_str("").unwrap(); - assert_eq!(options, PendingRowsBatcherOptions::default()); + let options: BatcherOptions = toml::from_str("").unwrap(); + assert_eq!(options, BatcherOptions::default()); assert!(!options.pending_rows_batching_enabled()); assert_eq!(options.max_batch_rows, 100_000); assert_eq!(options.max_concurrent_flushes, 256); - assert_eq!(options.worker_channel_capacity, 65526); + assert_eq!(options.worker_channel_capacity, 65_536); assert_eq!(options.max_inflight_requests, 3000); let serialized = toml::to_string(&options).unwrap(); assert_eq!( options, - toml::from_str::(&serialized).unwrap() + toml::from_str::(&serialized).unwrap() ); } + #[test] + fn test_worker_capacity_override() { + let options: FrontendOptions = toml::from_str( + r#" +[pending_rows_batcher] +worker_channel_capacity = 12345 +[pending_rows_batcher.logical_table] +worker_channel_capacity = 12345 +[prom_store] +enable = true +with_metric_engine = true +worker_channel_capacity = 12345 +"#, + ) + .unwrap(); + assert_eq!( + options.pending_rows_batcher.table.worker_channel_capacity, + 12_345 + ); + assert_eq!( + options + .pending_rows_batcher + .logical_table + .unwrap() + .worker_channel_capacity, + 12_345 + ); + assert_eq!(options.prom_store.worker_channel_capacity, 12_345); + } + #[test] fn test_partial_options_and_zero_controls() { - let options: PendingRowsBatcherOptions = + let options: BatcherOptions = toml::from_str("pending_rows_flush_interval = '5ms'").unwrap(); assert_eq!( options.pending_rows_flush_interval, Duration::from_millis(5) ); assert!(!options.pending_rows_batching_enabled()); - let enabled: PendingRowsBatcherOptions = + let enabled: BatcherOptions = toml::from_str("protocols = ['http_sql']\npending_rows_flush_interval = '5ms'") .unwrap(); assert!(enabled.pending_rows_batching_enabled()); @@ -140,7 +342,7 @@ mod tests { "worker_channel_capacity", "max_inflight_requests", ] { - let options: PendingRowsBatcherOptions = toml::from_str(&format!( + let options: BatcherOptions = toml::from_str(&format!( "protocols = ['influxdb']\npending_rows_flush_interval = '5ms'\n{field} = 0" )) .unwrap(); diff --git a/src/frontend/src/service_config/prom_store.rs b/src/frontend/src/service_config/prom_store.rs index f8429a86925..575ec600c15 100644 --- a/src/frontend/src/service_config/prom_store.rs +++ b/src/frontend/src/service_config/prom_store.rs @@ -27,9 +27,6 @@ pub struct PromStoreOptions { /// Validation mode while decoding Prometheus remote write requests. #[serde(default)] pub prom_validation_mode: PromValidationMode, - /// Enables experimental Prometheus remote write v2 native histogram ingestion. - #[serde(default)] - pub experimental_enable_prometheus_native_histogram: bool, #[serde(default, with = "humantime_serde")] pub pending_rows_flush_interval: Duration, #[serde(default = "default_max_batch_rows")] @@ -54,7 +51,7 @@ fn default_max_concurrent_flushes() -> usize { } fn default_worker_channel_capacity() -> usize { - 65526 + 65_536 } fn default_max_inflight_requests() -> usize { @@ -89,7 +86,6 @@ impl Default for PromStoreOptions { enable: true, with_metric_engine: true, prom_validation_mode: PromValidationMode::Strict, - experimental_enable_prometheus_native_histogram: false, pending_rows_flush_interval: Duration::ZERO, max_batch_rows: default_max_batch_rows(), max_concurrent_flushes: default_max_concurrent_flushes(), @@ -108,7 +104,6 @@ mod tests { use crate::service_config::prom_store::{ default_flow_notification_queue_capacity, default_max_batch_rows, default_max_concurrent_flushes, default_max_inflight_requests, - default_worker_channel_capacity, }; #[test] @@ -153,17 +148,13 @@ mod tests { assert!(default.enable); assert!(default.with_metric_engine); assert_eq!(default.prom_validation_mode, PromValidationMode::Strict); - assert!(!default.experimental_enable_prometheus_native_histogram); assert_eq!(default.pending_rows_flush_interval, Duration::ZERO); assert_eq!(default.max_batch_rows, default_max_batch_rows()); assert_eq!( default.max_concurrent_flushes, default_max_concurrent_flushes() ); - assert_eq!( - default.worker_channel_capacity, - default_worker_channel_capacity() - ); + assert_eq!(default.worker_channel_capacity, 65_536); assert_eq!( default.max_inflight_requests, default_max_inflight_requests() diff --git a/src/index/Cargo.toml b/src/index/Cargo.toml index b8518112610..2458a6591af 100644 --- a/src/index/Cargo.toml +++ b/src/index/Cargo.toml @@ -7,10 +7,8 @@ license.workspace = true [lints] workspace = true -[features] -vector_index = ["dep:nalgebra", "dep:usearch"] - [dependencies] +ahash.workspace = true async-trait.workspace = true asynchronous-codec = "0.7.0" bytemuck.workspace = true @@ -20,16 +18,14 @@ common-error.workspace = true common-macro.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -datatypes.workspace = true fastbloom = "0.8" fst.workspace = true futures.workspace = true greptime-proto.workspace = true itertools.workspace = true -jieba-rs = "0.10" +jieba-rs = "0.11" lazy_static.workspace = true mockall.workspace = true -nalgebra = { workspace = true, optional = true } pin-project.workspace = true prost.workspace = true puffin.workspace = true @@ -41,10 +37,9 @@ serde_json.workspace = true snafu.workspace = true store-api.workspace = true tantivy = { version = "0.26", features = ["zstd-compression"] } -tantivy-jieba = "0.20" +tantivy-jieba = "0.21" tokio.workspace = true tokio-util.workspace = true -usearch = { version = "2.21", default-features = false, features = ["fp16lib"], optional = true } uuid.workspace = true [dev-dependencies] @@ -63,3 +58,7 @@ harness = false [[bench]] name = "bytes_to_u64_vec" harness = false + +[[bench]] +name = "index_build_bench" +harness = false diff --git a/src/index/benches/index_build_bench.rs b/src/index/benches/index_build_bench.rs new file mode 100644 index 00000000000..4c6c3fac020 --- /dev/null +++ b/src/index/benches/index_build_bench.rs @@ -0,0 +1,424 @@ +// 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. + +//! Index build and bloom search benchmarks over synthetic observability data. + +use std::collections::{BTreeSet, HashMap}; +use std::hint::black_box; +use std::num::NonZeroUsize; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; + +use async_trait::async_trait; +use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main}; +use futures::{AsyncRead, AsyncReadExt}; +use index::bitmap::BitmapType; +use index::bloom_filter::applier::{BloomFilterApplier, InListPredicate}; +use index::bloom_filter::creator::BloomFilterCreator; +use index::bloom_filter::reader::BloomFilterReaderImpl; +use index::external_provider::{ExternalTempFileProvider, Reader, Writer}; +use index::fulltext_index::Config; +use index::fulltext_index::create::{BloomFilterFulltextIndexCreator, FulltextIndexCreator}; +use index::fulltext_index::tokenizer::{Analyzer, EnglishTokenizer}; +use index::inverted_index::create::InvertedIndexCreator; +use index::inverted_index::create::sort::external_sort::ExternalSorter; +use index::inverted_index::create::sort_create::SortIndexCreator; +use index::inverted_index::format::writer::InvertedIndexBlobWriter; +use puffin::puffin_manager::{PuffinWriter, PutOptions}; +use rand::seq::IndexedRandom; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +const ROWS: usize = 100_000; +const SEGMENT_ROWS: usize = 10240; +const ROW_GROUP_ROWS: usize = 102400; + +/// Drains blobs so the bloom fulltext creator can finish without a puffin file. +struct DrainPuffinWriter; + +#[async_trait] +impl PuffinWriter for DrainPuffinWriter { + async fn put_blob( + &mut self, + _key: &str, + raw_data: R, + _options: PutOptions, + _properties: HashMap, + ) -> puffin::error::Result + where + R: AsyncRead + Send, + { + let mut buf = Vec::new(); + Box::pin(raw_data).read_to_end(&mut buf).await.unwrap(); + Ok(buf.len() as u64) + } + + async fn put_dir( + &mut self, + _key: &str, + _dir: PathBuf, + _options: PutOptions, + _properties: HashMap, + ) -> puffin::error::Result { + unreachable!("bloom fulltext index only writes blobs") + } + + fn set_footer_lz4_compressed(&mut self, _lz4_compressed: bool) {} + + async fn finish(self) -> puffin::error::Result { + Ok(0) + } +} + +/// The benchmarks set no memory limit, so nothing is spilled. +struct NoSpill; + +#[async_trait] +impl ExternalTempFileProvider for NoSpill { + async fn create(&self, _: &str, _: &str) -> Result { + unreachable!("no memory limit is set") + } + + async fn read_all(&self, _: &str) -> Result, index::error::Error> { + Ok(vec![]) + } +} + +fn uuid(rng: &mut ChaCha8Rng) -> String { + let h = format!("{:032x}", rng.random::()); + format!( + "{}-{}-{}-{}-{}", + &h[0..8], + &h[8..12], + &h[12..16], + &h[16..20], + &h[20..32] + ) +} + +/// Access, auth, timeout and error log lines with request ids, IPs and numbers: most +/// tokens of a segment are distinct. +fn service_logs(rows: usize) -> Vec { + let mut rng = ChaCha8Rng::seed_from_u64(1); + let paths = ["users", "orders", "items", "carts", "payments", "sessions"]; + (0..rows) + .map(|_| { + let ip = format!( + "10.{}.{}.{}", + rng.random_range(0..16), + rng.random_range(0..256), + rng.random_range(1..255) + ); + match rng.random_range(0..10) { + 0..=5 => format!( + "INFO GET /api/v1/{}/{}/detail?page={} 200 {}ms request_id={} client={}", + paths.choose(&mut rng).unwrap(), + rng.random_range(0..200000), + rng.random_range(0..50), + rng.random_range(1..2000), + uuid(&mut rng), + ip + ), + 6..=7 => format!( + "WARN connection to {}:{} timed out after {}ms retry={}", + ip, + rng.random_range(1024..65535), + rng.random_range(100..30000), + rng.random_range(0..5) + ), + _ => format!( + "ERROR failed to process message {} from queue orders-{}: deadline exceeded", + uuid(&mut rng), + rng.random_range(0..64) + ), + } + }) + .collect() +} + +/// Java stack traces: long lines where a few frame names repeat many times. +fn stack_trace_logs(rows: usize) -> Vec { + let mut rng = ChaCha8Rng::seed_from_u64(2); + let frames = [ + "at org.apache.kafka.clients.consumer.KafkaConsumer.poll(KafkaConsumer.java:1250)", + "at com.example.orders.OrderService.process(OrderService.java:88)", + "at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1136)", + "at java.base/java.lang.Thread.run(Thread.java:833)", + ]; + (0..rows) + .map(|_| { + let mut line = format!( + "ERROR request {} failed: java.lang.IllegalStateException: retry exhausted", + uuid(&mut rng) + ); + for _ in 0..40 { + line.push(' '); + line.push_str(frames.choose(&mut rng).unwrap()); + } + line + }) + .collect() +} + +/// Trace ids with temporal locality: each trace spreads over about a hundred rows. +fn trace_ids(rows: usize) -> Vec { + let mut rng = ChaCha8Rng::seed_from_u64(3); + let traces = (0..rows / 10) + .map(|_| format!("{:032x}", rng.random::())) + .collect::>(); + (0..rows) + .map(|i| { + let j = (i / 10) as i64 + rng.random_range(-50i64..=50); + traces[j.clamp(0, traces.len() as i64 - 1) as usize].clone() + }) + .collect() +} + +/// Prometheus-like series sorted by labels; each series has `samples` rows. +fn series(num_series: usize) -> Vec<[String; 6]> { + let mut rng = ChaCha8Rng::seed_from_u64(4); + let mut series = (0..num_series) + .map(|s| { + let pod = s / 2; + let ns = pod % 400; + [ + format!("region-{}", ns % 4), + format!("az-{}", ns % 12), + format!("cluster-{}", ns % 40), + format!("namespace-{ns}"), + format!("pod-{:05x}-{pod}", rng.random_range(0..0xfffff)), + format!("container-{}", s % 2), + ] + }) + .collect::>(); + series.sort(); + series +} + +fn bloom_creator() -> BloomFilterCreator { + BloomFilterCreator::new( + SEGMENT_ROWS, + 0.01, + Arc::new(NoSpill), + Arc::new(AtomicUsize::new(0)), + None, + ) +} + +fn bench_bloom_fulltext_build(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bloom_fulltext_build"); + group.sample_size(10); + + for (name, lines) in [ + ("service_logs", service_logs(ROWS)), + ("stack_traces", stack_trace_logs(ROWS / 10)), + ] { + group.throughput(Throughput::Elements(lines.len() as u64)); + group.bench_function(name, |b| { + b.iter_batched( + || { + BloomFilterFulltextIndexCreator::new( + Config::default(), + SEGMENT_ROWS, + 0.01, + Arc::new(NoSpill), + Arc::new(AtomicUsize::new(0)), + None, + ) + }, + |mut creator| { + let lines = &lines; + rt.block_on(async move { + for line in lines { + creator.push_text(line).await.unwrap(); + } + let written = creator + .finish(&mut DrainPuffinWriter, "blob", PutOptions::default()) + .await + .unwrap(); + black_box(written) + }) + }, + BatchSize::LargeInput, + ) + }); + } + group.finish(); +} + +fn bench_bloom_skipping_build(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bloom_skipping_build"); + group.sample_size(10); + + let trace_ids = trace_ids(ROWS * 10); + group.throughput(Throughput::Elements(trace_ids.len() as u64)); + group.bench_function("trace_id", |b| { + b.iter_batched( + bloom_creator, + |mut creator| { + let trace_ids = &trace_ids; + rt.block_on(async move { + for id in trace_ids { + creator + .push_n_row_elem(1, Some(id.as_bytes())) + .await + .unwrap(); + } + let mut out = Vec::new(); + creator.finish(&mut out).await.unwrap(); + black_box(out) + }) + }, + BatchSize::LargeInput, + ) + }); + + // A sorted tag pushed as one run per series, as the flat SST path does. + let series = series(ROWS); + let samples = 10; + group.throughput(Throughput::Elements((series.len() * samples) as u64)); + group.bench_function("sorted_tag_runs", |b| { + b.iter_batched( + bloom_creator, + |mut creator| { + let series = &series; + rt.block_on(async move { + for s in series { + creator + .push_n_row_elem(samples, Some(s[4].as_bytes())) + .await + .unwrap(); + } + let mut out = Vec::new(); + creator.finish(&mut out).await.unwrap(); + black_box(out) + }) + }, + BatchSize::LargeInput, + ) + }); + group.finish(); +} + +fn bench_inverted_build(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("inverted_build"); + group.sample_size(10); + + let series = series(ROWS); + let samples = 10; + let names = ["region", "az", "cluster", "namespace", "pod", "container"]; + group.throughput(Throughput::Elements((series.len() * samples) as u64)); + group.bench_function("six_tags", |b| { + b.iter_batched( + || { + let factory = ExternalSorter::factory( + Arc::new(NoSpill), + None, + Arc::new(AtomicUsize::new(0)), + None, + ); + SortIndexCreator::new(factory, NonZeroUsize::new(1024).unwrap()) + }, + |mut creator| { + let series = &series; + rt.block_on(async move { + for s in series { + for (name, value) in names.iter().zip(s) { + let spill = + creator.push_with_name_n(name, Some(value.as_bytes()), samples); + assert!(!spill, "no memory limit is set"); + } + } + let mut out = Vec::new(); + let mut writer = InvertedIndexBlobWriter::new(&mut out); + creator + .finish(&mut writer, BitmapType::Roaring) + .await + .unwrap(); + black_box(out) + }) + }, + BatchSize::LargeInput, + ) + }); + group.finish(); +} + +#[allow(clippy::single_range_in_vec_init)] +fn bench_bloom_search(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bloom_search"); + + let lines = service_logs(ROWS * 10); + let analyzer = Analyzer::new(Box::new(EnglishTokenizer), false); + let blob = rt.block_on(async { + let mut creator = bloom_creator(); + for line in &lines { + creator + .push_row_elems(analyzer.analyze_text(line).unwrap()) + .await + .unwrap(); + } + let mut out = Vec::new(); + creator.finish(&mut out).await.unwrap(); + out + }); + + let num_rows = lines.len(); + // Each row group tests every segment's filter against the three probes. + let predicates = ["timed", "deadline", "nonexistent"] + .iter() + .map(|term| InListPredicate { + list: BTreeSet::from([term.as_bytes().to_vec()]), + }) + .collect::>(); + group.bench_function("three_terms_per_row_group", |b| { + b.iter_batched( + || { + let reader = BloomFilterReaderImpl::new(blob.clone()); + rt.block_on(BloomFilterApplier::new(Box::new(reader))) + .unwrap() + }, + |mut applier| { + rt.block_on(async { + let mut matched = 0; + for start in (0..num_rows).step_by(ROW_GROUP_ROWS) { + let end = (start + ROW_GROUP_ROWS).min(num_rows); + matched += applier + .search(&predicates, &[start..end], None) + .await + .unwrap() + .len(); + } + black_box(matched) + }) + }, + BatchSize::SmallInput, + ) + }); + group.finish(); +} + +criterion_group!( + benches, + bench_bloom_fulltext_build, + bench_bloom_skipping_build, + bench_inverted_build, + bench_bloom_search +); +criterion_main!(benches); diff --git a/src/index/src/bloom_filter.rs b/src/index/src/bloom_filter.rs index eb818a0f5a0..156be09e69e 100644 --- a/src/index/src/bloom_filter.rs +++ b/src/index/src/bloom_filter.rs @@ -17,5 +17,73 @@ pub mod creator; pub mod error; pub mod reader; +use std::hash::{BuildHasher, BuildHasherDefault, Hasher}; +use std::sync::LazyLock; + /// The seed used for the Bloom filter. pub const SEED: u128 = 42; + +static ELEMENT_HASHER: LazyLock = + LazyLock::new(|| fastbloom::DefaultHasher::seeded(&SEED.to_be_bytes())); + +/// Returns the hash fastbloom derives for `elem` with the persisted [`SEED`]. +/// +/// A filter built from these hashes with [`PrehashedBuildHasher`] has exactly the same +/// bits as one built by inserting the elements themselves, so files stay readable by +/// both paths. +pub fn element_hash(elem: &[u8]) -> u64 { + ELEMENT_HASHER.hash_one(elem) +} + +/// Hasher that returns an already computed [`element_hash`] unchanged. +#[derive(Default)] +pub struct PrehashedHasher(u64); + +impl Hasher for PrehashedHasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, _bytes: &[u8]) { + unreachable!("PrehashedHasher only accepts u64 hashes") + } + + fn write_u64(&mut self, hash: u64) { + self.0 = hash; + } +} + +pub type PrehashedBuildHasher = BuildHasherDefault; + +/// A persisted bloom filter probed by [`element_hash`] values. +pub type PrehashedBloomFilter = fastbloom::BloomFilter<512, PrehashedBuildHasher>; + +#[cfg(test)] +mod tests { + use fastbloom::BloomFilter; + + use super::*; + + #[test] + fn test_prehashed_filter_matches_seeded_filter() { + // Persisted filters are read back with `.seed(&SEED)`, so building them from + // precomputed hashes must set exactly the same bits. + for count in [0usize, 1, 7, 100, 5000] { + let elems = (0..count) + .map(|i| format!("elem-{i}").into_bytes()) + .chain([Vec::new()]) + .collect::>(); + let mut seeded = BloomFilter::with_false_pos(0.01) + .seed(&SEED) + .expected_items(elems.len()); + let mut prehashed = BloomFilter::with_false_pos(0.01) + .hasher(PrehashedBuildHasher::default()) + .expected_items(elems.len()); + for elem in &elems { + seeded.insert(elem); + prehashed.insert(&element_hash(elem)); + } + assert_eq!(seeded.as_slice(), prehashed.as_slice()); + } + } +} diff --git a/src/index/src/bloom_filter/applier.rs b/src/index/src/bloom_filter/applier.rs index db219c9e611..96a1c6ff17f 100644 --- a/src/index/src/bloom_filter/applier.rs +++ b/src/index/src/bloom_filter/applier.rs @@ -14,14 +14,20 @@ use std::collections::BTreeSet; use std::ops::Range; +use std::sync::Arc; -use fastbloom::BloomFilter; use greptime_proto::v1::index::BloomFilterMeta; use itertools::Itertools; use crate::Bytes; use crate::bloom_filter::error::Result; use crate::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader}; +use crate::bloom_filter::{PrehashedBloomFilter, element_hash}; + +/// Filter bytes one batch of [`BloomFilterApplier::search_groups`] reads. A single row +/// group larger than this is still read as one batch, so this is not a memory limit; a +/// batch holds both its raw bytes and the decoded filters. +const MAX_BATCH_FILTER_BYTES: u64 = 8 * 1024 * 1024; /// `InListPredicate` contains a list of acceptable values. A value needs to match at least /// one of the elements (logical OR semantic) for the predicate to be satisfied. @@ -33,7 +39,7 @@ pub struct InListPredicate { pub struct BloomFilterApplier { reader: Box, - meta: BloomFilterMeta, + meta: Arc, } impl BloomFilterApplier { @@ -43,6 +49,93 @@ impl BloomFilterApplier { Ok(Self { reader, meta }) } + /// Runs [`Self::search`] over several groups of ranges (e.g. row groups) and keeps in + /// each group only its matching ranges. + /// + /// Consecutive groups are searched together, one read per batch, as long as the + /// filters of a batch stay within [`MAX_BATCH_FILTER_BYTES`]. This saves a round trip + /// per group on object storage without holding the filters of a whole file at once. + /// Groups must be ordered and their ranges sorted and disjoint, as for `search`. + pub async fn search_groups( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result<()> { + self.search_groups_in_batches(predicates, groups, metrics, MAX_BATCH_FILTER_BYTES) + .await + } + + async fn search_groups_in_batches( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + mut metrics: Option<&mut BloomFilterReadMetrics>, + max_batch_bytes: u64, + ) -> Result<()> { + let mut start = 0; + while start < groups.len() { + // Mirrors `load_bloom_filters`: segments map to filter locations in order and + // only consecutive equal locations share a read, so this is the exact number of + // bytes the batch requests. + let mut last_loc = None; + let mut batch_bytes = 0; + let mut end = start; + while end < groups.len() { + let mut group_last = last_loc; + let mut bytes = 0; + for seg in self.row_ranges_to_segments(groups[end]) { + let loc = self.meta.segment_loc_indices[seg]; + if group_last != Some(loc) { + bytes += self.meta.bloom_filter_locs[loc as usize].size; + group_last = Some(loc); + } + } + // A group whose filters alone exceed the budget still forms its own batch. + if end > start && batch_bytes + bytes > max_batch_bytes { + break; + } + last_loc = group_last; + batch_bytes += bytes; + end += 1; + } + self.search_batch(predicates, &mut groups[start..end], metrics.as_deref_mut()) + .await?; + start = end; + } + Ok(()) + } + + /// Searches `groups` with one `search`, i.e. one read for all their filters. + async fn search_batch( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result<()> { + let all = groups + .iter() + .flat_map(|g| g.iter().cloned()) + .collect::>(); + if all.is_empty() { + return Ok(()); + } + // Each matched range lies within one input range, so it belongs to one group. + let mut matched = self + .search(predicates, &all, metrics) + .await? + .into_iter() + .peekable(); + for group in groups.iter_mut() { + for range in std::mem::take(*group) { + while let Some(m) = matched.next_if(|m| m.start < range.end) { + group.push(m); + } + } + } + Ok(()) + } + /// Searches ranges of rows that match all the given predicates in the search ranges. /// Each predicate represents an OR condition of probes, and all predicates must match (AND semantics). /// The logic is: (probe1 OR probe2 OR ...) AND (probe3 OR probe4 OR ...) @@ -97,7 +190,7 @@ impl BloomFilterApplier { &mut self, segments: &[usize], metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result<(Vec<(u64, usize)>, Vec)> { + ) -> Result<(Vec<(u64, usize)>, Vec)> { let segment_locations = segments .iter() .map(|&seg| (self.meta.segment_loc_indices[seg], seg)) @@ -122,11 +215,15 @@ impl BloomFilterApplier { fn find_matching_rows( &self, segment_locations: Vec<(u64, usize)>, - bloom_filters: Vec, + bloom_filters: Vec, predicates: &[InListPredicate], ) -> Vec> { let rows_per_segment = self.meta.rows_per_segment as usize; let mut matching_row_ranges = Vec::with_capacity(bloom_filters.len()); + let predicate_hashes = predicates + .iter() + .map(|p| p.list.iter().map(|v| element_hash(v)).collect::>()) + .collect::>(); // Group segments by their location index (since they have the same bloom filter) and check if they match all predicates for ((_loc_index, group), bloom_filter) in segment_locations @@ -136,12 +233,9 @@ impl BloomFilterApplier { .zip(bloom_filters.iter()) { // Check if this bloom filter matches each predicate (AND semantics) - let matches_all_predicates = predicates.iter().all(|predicate| { + let matches_all_predicates = predicate_hashes.iter().all(|hashes| { // For each predicate, at least one probe must match (OR semantics) - predicate - .list - .iter() - .any(|probe| bloom_filter.contains(probe)) + hashes.iter().any(|hash| bloom_filter.contains(hash)) }); if !matches_all_predicates { @@ -217,6 +311,222 @@ mod tests { use crate::bloom_filter::reader::BloomFilterReaderImpl; use crate::external_provider::MockExternalTempFileProvider; + #[tokio::test] + async fn test_search_groups_matches_per_group_search() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + // Row i holds "v{i / 3}", so values straddle segment boundaries. + for i in 0..40 { + creator + .push_row_elems([format!("v{}", i / 3).into_bytes()]) + .await + .unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let bytes = writer.into_inner(); + + // Row groups of 10 rows, some already narrowed by other predicates. + let groups = vec![ + vec![0..10], + vec![10..13, 15..20], + vec![], + vec![30..33, 37..40], + ]; + for values in [ + vec!["v1"], + vec!["v4", "v5"], + vec!["v3", "v10", "v12"], + vec!["x"], + ] { + let predicates = vec![InListPredicate { + list: values.iter().map(|v| v.as_bytes().to_vec()).collect(), + }]; + let mut applier = + BloomFilterApplier::new(Box::new(BloomFilterReaderImpl::new(bytes.clone()))) + .await + .unwrap(); + let filter_size = applier.meta.bloom_filter_locs[0].size; + let mut expected = Vec::new(); + for group in &groups { + expected.push(if group.is_empty() { + vec![] + } else { + applier.search(&predicates, group, None).await.unwrap() + }); + } + // Budgets of zero (a batch per group), one filter and unlimited (one batch). + for budget in [0, filter_size, u64::MAX] { + let mut actual = groups.clone(); + let mut refs = actual.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, budget) + .await + .unwrap(); + assert_eq!(actual, expected, "values: {values:?}, budget: {budget}"); + } + } + } + + /// Records the bytes of every `read_vec`, i.e. of every batch. + struct RecordingReader { + inner: BloomFilterReaderImpl>, + reads: Arc>>, + } + + #[async_trait::async_trait] + impl BloomFilterReader for RecordingReader { + async fn range_read( + &self, + offset: u64, + size: u32, + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result { + self.inner.range_read(offset, size, metrics).await + } + + async fn read_vec( + &self, + ranges: &[Range], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result> { + let bytes = ranges.iter().map(|r| r.end - r.start).sum(); + self.reads.lock().unwrap().push(bytes); + self.inner.read_vec(ranges, metrics).await + } + + async fn metadata( + &self, + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result> { + self.inner.metadata(metrics).await + } + } + + #[tokio::test] + #[allow(clippy::single_range_in_vec_init)] + async fn test_search_groups_respects_budget() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + // Distinct values everywhere: 10 segments with filters of the same size. + for i in 0..40 { + creator + .push_row_elems([format!("v{i}").into_bytes()]) + .await + .unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let bytes = writer.into_inner(); + let predicates = vec![InListPredicate { + list: BTreeSet::from([b"v1".to_vec()]), + }]; + + let reads = Arc::new(std::sync::Mutex::new(Vec::new())); + let reader = RecordingReader { + inner: BloomFilterReaderImpl::new(bytes), + reads: reads.clone(), + }; + let mut applier = BloomFilterApplier::new(Box::new(reader)).await.unwrap(); + let filter_size = applier.meta.bloom_filter_locs[0].size; + + // Row groups of 6 rows: most share a boundary segment with their neighbor, which + // a batch reads only once. + let groups = (0..40) + .step_by(6) + .map(|s| vec![s..(s + 6).min(40)]) + .collect::>(); + for (budget, expected_reads) in [ + (u64::MAX, vec![10 * filter_size]), + // Batches close before exceeding the budget; shared boundary filters count once. + ( + 5 * filter_size, + vec![5 * filter_size, 5 * filter_size, filter_size], + ), + // Smaller than any row group: one row group per batch, still above the budget. + ( + filter_size, + vec![ + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + filter_size, + ], + ), + ] { + reads.lock().unwrap().clear(); + let mut actual = groups.clone(); + let mut refs = actual.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, budget) + .await + .unwrap(); + assert_eq!(*reads.lock().unwrap(), expected_reads, "budget: {budget}"); + } + + // Segments 0..3 hold the same value and share one filter, which a batch counts + // and reads once even across row groups. + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + for i in 0..24 { + let value = if i < 12 { + "a".to_string() + } else { + format!("v{i}") + }; + creator.push_row_elems([value.into_bytes()]).await.unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let reader = RecordingReader { + inner: BloomFilterReaderImpl::new(writer.into_inner()), + reads: reads.clone(), + }; + let mut applier = BloomFilterApplier::new(Box::new(reader)).await.unwrap(); + assert_eq!(applier.meta.bloom_filter_locs.len(), 4); + let filter_size = applier.meta.bloom_filter_locs[0].size; + assert!( + applier + .meta + .bloom_filter_locs + .iter() + .all(|l| l.size == filter_size) + ); + reads.lock().unwrap().clear(); + let mut groups = (0..24) + .step_by(6) + .map(|s| vec![s..s + 6]) + .collect::>(); + let mut refs = groups.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, 2 * filter_size) + .await + .unwrap(); + // Row groups 0 and 1 need only the shared filter; 2 and 3 need two each. + assert_eq!( + *reads.lock().unwrap(), + vec![filter_size, 2 * filter_size, 2 * filter_size] + ); + } + #[tokio::test] #[allow(clippy::single_range_in_vec_init)] async fn test_appliter() { diff --git a/src/index/src/bloom_filter/creator.rs b/src/index/src/bloom_filter/creator.rs index df23ca2352a..dc1dc54d289 100644 --- a/src/index/src/bloom_filter/creator.rs +++ b/src/index/src/bloom_filter/creator.rs @@ -26,8 +26,8 @@ use prost::Message; use snafu::ResultExt; use crate::Bytes; -use crate::bloom_filter::SEED; use crate::bloom_filter::error::{IoSnafu, Result}; +use crate::bloom_filter::{PrehashedBuildHasher, element_hash}; use crate::external_provider::ExternalTempFileProvider; /// `BloomFilterCreator` is responsible for creating and managing bloom filters @@ -52,8 +52,11 @@ pub struct BloomFilterCreator { /// Row count that added to the bloom filter so far. accumulated_row_count: usize, - /// A set of distinct elements in the current segment. - cur_seg_distinct_elems: HashSet, + /// Distinct element hashes (see [`element_hash`]) in the current segment. + /// + /// Elements with equal hashes set the same bits, so deduplicating by hash loses + /// nothing and avoids copying the values. + cur_seg_distinct_elems: HashSet, /// The memory usage of the current segment's distinct elements. cur_seg_distinct_elems_mem_usage: usize, @@ -106,98 +109,36 @@ impl BloomFilterCreator { /// reaches `rows_per_segment`, it finalizes the current segment. pub async fn push_n_row_elems( &mut self, - mut nrows: usize, + nrows: usize, elems: impl IntoIterator, ) -> Result<()> { - if nrows == 0 { - return Ok(()); - } - if nrows == 1 { - return self.push_row_elems(elems).await; - } - - let elems = elems.into_iter().collect::>(); - while nrows > 0 { - let rows_to_seg_end = - self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment); - let rows_to_push = nrows.min(rows_to_seg_end); - nrows -= rows_to_push; - - self.accumulated_row_count += rows_to_push; - - let mut mem_diff = 0; - for elem in &elems { - let len = elem.len(); - let is_new = self.cur_seg_distinct_elems.insert(elem.clone()); - if is_new { - mem_diff += len; - } - } - self.cur_seg_distinct_elems_mem_usage += mem_diff; - self.global_memory_usage - .fetch_add(mem_diff, Ordering::Relaxed); - - if self - .accumulated_row_count - .is_multiple_of(self.rows_per_segment) - { - self.finalize_segment().await?; - self.finalized_row_count = self.accumulated_row_count; - } - } - - Ok(()) + let hashes = elems + .into_iter() + .map(|e| element_hash(&e)) + .collect::>(); + self.push_n_row_hashes(nrows, &hashes).await } - /// Adds `nrows` copies of a single borrowed value (or null), copying it only when - /// it is new to a segment. Row counts advance for nulls as well. - pub async fn push_n_row_elem(&mut self, mut nrows: usize, elem: Option<&[u8]>) -> Result<()> { - while nrows > 0 { - let rows_to_seg_end = - self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment); - let rows_to_push = nrows.min(rows_to_seg_end); - nrows -= rows_to_push; - self.accumulated_row_count += rows_to_push; - - if let Some(elem) = elem { - let old_len = self.cur_seg_distinct_elems.len(); - // A borrowed entry lookup avoids hashing new values twice and only - // allocates when the value is absent from the current segment. - self.cur_seg_distinct_elems - .get_or_insert_with(elem, <[u8]>::to_vec); - if self.cur_seg_distinct_elems.len() != old_len { - self.cur_seg_distinct_elems_mem_usage += elem.len(); - self.global_memory_usage - .fetch_add(elem.len(), Ordering::Relaxed); - } - } - if self - .accumulated_row_count - .is_multiple_of(self.rows_per_segment) - { - self.finalize_segment().await?; - self.finalized_row_count = self.accumulated_row_count; - } + /// Adds `nrows` copies of a single borrowed value (or null). Row counts advance for + /// nulls as well. + pub async fn push_n_row_elem(&mut self, nrows: usize, elem: Option<&[u8]>) -> Result<()> { + match elem { + Some(elem) => self.push_n_row_hashes(nrows, &[element_hash(elem)]).await, + None => self.push_n_row_hashes(nrows, &[]).await, } - Ok(()) } /// Adds a row of elements to the bloom filter. If the number of accumulated rows /// reaches `rows_per_segment`, it finalizes the current segment. pub async fn push_row_elems(&mut self, elems: impl IntoIterator) -> Result<()> { - self.accumulated_row_count += 1; + self.push_row_hashes(elems.into_iter().map(|e| element_hash(&e))) + .await + } - let mut mem_diff = 0; - for elem in elems.into_iter() { - let len = elem.len(); - let is_new = self.cur_seg_distinct_elems.insert(elem); - if is_new { - mem_diff += len; - } - } - self.cur_seg_distinct_elems_mem_usage += mem_diff; - self.global_memory_usage - .fetch_add(mem_diff, Ordering::Relaxed); + /// Adds a row of element hashes computed by [`element_hash`]. + pub async fn push_row_hashes(&mut self, hashes: impl IntoIterator) -> Result<()> { + self.accumulated_row_count += 1; + self.insert_hashes(hashes); if self .accumulated_row_count @@ -210,6 +151,43 @@ impl BloomFilterCreator { Ok(()) } + /// Adds `nrows` rows that all contain the element hashes in `hashes`. + pub async fn push_n_row_hashes(&mut self, mut nrows: usize, hashes: &[u64]) -> Result<()> { + while nrows > 0 { + let rows_to_seg_end = + self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment); + let rows_to_push = nrows.min(rows_to_seg_end); + nrows -= rows_to_push; + self.accumulated_row_count += rows_to_push; + self.insert_hashes(hashes.iter().copied()); + + if self + .accumulated_row_count + .is_multiple_of(self.rows_per_segment) + { + self.finalize_segment().await?; + self.finalized_row_count = self.accumulated_row_count; + } + } + + Ok(()) + } + + fn insert_hashes(&mut self, hashes: impl IntoIterator) { + let old_len = self.cur_seg_distinct_elems.len(); + // Not `extend`: it reserves for the iterator's length, which counts duplicate + // tokens, and the capacity survives `drain` at segment boundaries. + for hash in hashes { + self.cur_seg_distinct_elems.insert(hash); + } + let mem_diff = (self.cur_seg_distinct_elems.len() - old_len) * size_of::(); + if mem_diff > 0 { + self.cur_seg_distinct_elems_mem_usage += mem_diff; + self.global_memory_usage + .fetch_add(mem_diff, Ordering::Relaxed); + } + } + /// Finalizes any remaining segments and writes the bloom filters and metadata to the provided writer. pub async fn finish(&mut self, mut writer: impl AsyncWrite + Unpin) -> Result<()> { if self.accumulated_row_count > self.finalized_row_count { @@ -286,16 +264,36 @@ mod tests { use futures::io::Cursor; use super::*; + use crate::bloom_filter::SEED; use crate::external_provider::MockExternalTempFileProvider; /// Converts a slice of bytes to a vector of `u64`. pub fn u64_vec_from_bytes(bytes: &[u8]) -> Vec { bytes - .chunks_exact(std::mem::size_of::()) - .map(|chunk| u64::from_le_bytes(chunk.try_into().unwrap())) + .as_chunks::<{ std::mem::size_of::() }>() + .0 + .iter() + .map(|chunk| u64::from_le_bytes(*chunk)) .collect() } + #[tokio::test] + async fn test_duplicate_hashes_do_not_grow_segment_set() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + creator + .push_row_hashes(std::iter::repeat_n(7, 1_000_000)) + .await + .unwrap(); + assert_eq!(creator.cur_seg_distinct_elems.len(), 1); + assert!(creator.cur_seg_distinct_elems.capacity() < 16); + } + #[tokio::test] async fn test_bloom_filter_creator() { let mut writer = Cursor::new(Vec::new()); diff --git a/src/index/src/bloom_filter/creator/finalize_segment.rs b/src/index/src/bloom_filter/creator/finalize_segment.rs index dd7d03864e4..4cbce66edce 100644 --- a/src/index/src/bloom_filter/creator/finalize_segment.rs +++ b/src/index/src/bloom_filter/creator/finalize_segment.rs @@ -22,8 +22,7 @@ use futures::stream::StreamExt; use futures::{AsyncWriteExt, Stream, stream}; use snafu::ResultExt; -use crate::Bytes; -use crate::bloom_filter::creator::SEED; +use crate::bloom_filter::PrehashedBuildHasher; use crate::bloom_filter::creator::intermediate_codec::IntermediateBloomFilterCodecV1; use crate::bloom_filter::error::{IntermediateSnafu, IoSnafu, Result}; use crate::external_provider::ExternalTempFileProvider; @@ -98,14 +97,14 @@ impl FinalizedBloomFilterStorage { /// If the memory usage exceeds the threshold, flushes the in-memory Bloom filters to disk. pub async fn add( &mut self, - elems: impl IntoIterator, + elem_hashes: impl IntoIterator, element_count: usize, ) -> Result<()> { let mut bf = BloomFilter::with_false_pos(self.false_positive_rate) - .seed(&SEED) + .hasher(PrehashedBuildHasher::default()) .expected_items(element_count); - for elem in elems.into_iter() { - bf.insert(&elem); + for hash in elem_hashes.into_iter() { + bf.insert(&hash); } let fbf = FinalizedBloomFilterSegment::from(bf, element_count); @@ -231,7 +230,7 @@ pub struct FinalizedBloomFilterSegment { } impl FinalizedBloomFilterSegment { - fn from(bf: BloomFilter, elem_count: usize) -> Self { + fn from(bf: BloomFilter<512, S>, elem_count: usize) -> Self { let bf_slice = bf.as_slice(); let mut bloom_filter_bytes = Vec::with_capacity(std::mem::size_of_val(bf_slice)); for &x in bf_slice { @@ -256,6 +255,7 @@ mod tests { use super::*; use crate::bloom_filter::creator::tests::u64_vec_from_bytes; + use crate::bloom_filter::{SEED, element_hash}; use crate::external_provider::MockExternalTempFileProvider; #[tokio::test] @@ -300,11 +300,12 @@ mod tests { let dup_batch = 200; for i in 0..(batch - dup_batch) { - let elems = (elem_count * i..elem_count * (i + 1)).map(|x| x.to_string().into_bytes()); + let elems = (elem_count * i..elem_count * (i + 1)) + .map(|x| element_hash(x.to_string().as_bytes())); storage.add(elems, elem_count).await.unwrap(); } for _ in 0..dup_batch { - storage.add(Some(vec![]), 1).await.unwrap(); + storage.add(Some(element_hash(&[])), 1).await.unwrap(); } // Flush happens. @@ -356,7 +357,7 @@ mod tests { let batch = 1000; for _ in 0..batch { - storage.add(Some(vec![]), 1).await.unwrap(); + storage.add(Some(element_hash(&[])), 1).await.unwrap(); } // Drain the storage. diff --git a/src/index/src/bloom_filter/reader.rs b/src/index/src/bloom_filter/reader.rs index a9e08694e74..15cf957b7b6 100644 --- a/src/index/src/bloom_filter/reader.rs +++ b/src/index/src/bloom_filter/reader.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::ops::{Range, Rem}; +use std::sync::Arc; use std::time::{Duration, Instant}; use async_trait::async_trait; @@ -24,10 +25,10 @@ use greptime_proto::v1::index::{BloomFilterLoc, BloomFilterMeta}; use prost::Message; use snafu::{ResultExt, ensure}; -use crate::bloom_filter::SEED; use crate::bloom_filter::error::{ DecodeProtoSnafu, FileSizeTooSmallSnafu, IoSnafu, Result, UnexpectedMetaSizeSnafu, }; +use crate::bloom_filter::{PrehashedBloomFilter, PrehashedBuildHasher, SEED}; /// Minimum size of the bloom filter, which is the size of the length of the bloom filter. const BLOOM_META_LEN_SIZE: u64 = 4; @@ -164,7 +165,7 @@ pub trait BloomFilterReader: Sync { async fn metadata( &self, metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result; + ) -> Result>; /// Reads a bloom filter with the given location. async fn bloom_filter( @@ -180,11 +181,12 @@ pub trait BloomFilterReader: Sync { Ok(bm) } + /// Reads multiple bloom filters; probe them with [`crate::bloom_filter::element_hash`]. async fn bloom_filter_vec( &self, locs: &[BloomFilterLoc], metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result> { + ) -> Result> { let ranges = locs .iter() .map(|l| l.offset..l.offset + l.size) @@ -195,7 +197,7 @@ pub trait BloomFilterReader: Sync { for (bs, loc) in bss.into_iter().zip(locs.iter()) { let vec = bytes_to_u64_vec(&bs); let bm = BloomFilter::from_vec(vec) - .seed(&SEED) + .hasher(PrehashedBuildHasher::default()) .expected_items(loc.element_count as _); result.push(bm); } @@ -265,13 +267,13 @@ impl BloomFilterReader for BloomFilterReaderImpl { async fn metadata( &self, metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result { + ) -> Result> { let metadata = self.reader.metadata().await.context(IoSnafu)?; let file_size = metadata.content_length; let mut meta_reader = BloomFilterMetaReader::new(&self.reader, file_size, Some(DEFAULT_PREFETCH_SIZE)); - meta_reader.metadata(metrics).await + meta_reader.metadata(metrics).await.map(Arc::new) } } diff --git a/src/index/src/fulltext_index/create/bloom_filter.rs b/src/index/src/fulltext_index/create/bloom_filter.rs index 28648f9b540..a2db31418eb 100644 --- a/src/index/src/fulltext_index/create/bloom_filter.rs +++ b/src/index/src/fulltext_index/create/bloom_filter.rs @@ -74,11 +74,12 @@ impl BloomFilterFulltextIndexCreator { #[async_trait] impl FulltextIndexCreator for BloomFilterFulltextIndexCreator { async fn push_text(&mut self, text: &str) -> Result<()> { - let tokens = self.analyzer.analyze_text(text)?; + let mut token_buf = Vec::new(); + let hashes = self.analyzer.analyze_text_hashes(text, &mut token_buf); self.inner .as_mut() .context(AbortedSnafu)? - .push_row_elems(tokens) + .push_row_hashes(hashes) .await .map_err(BoxedError::new) .context(ExternalSnafu)?; diff --git a/src/index/src/fulltext_index/tokenizer.rs b/src/index/src/fulltext_index/tokenizer.rs index 3afc826e6f1..0a1cc47aa8e 100644 --- a/src/index/src/fulltext_index/tokenizer.rs +++ b/src/index/src/fulltext_index/tokenizer.rs @@ -13,6 +13,7 @@ // limitations under the License. use crate::Bytes; +use crate::bloom_filter::element_hash; use crate::fulltext_index::error::Result; lazy_static::lazy_static! { @@ -140,6 +141,30 @@ impl Analyzer { } } + /// Returns the bloom filter hash of each token in the given text. + /// + /// Equivalent to hashing every token returned by [`Analyzer::analyze_text`] with + /// [`element_hash`]. Only case-insensitive non-ASCII tokens allocate, in `to_lowercase`; + /// case-insensitive ASCII tokens are lowercased in `buf`. + pub fn analyze_text_hashes<'a>( + &self, + text: &'a str, + buf: &'a mut Vec, + ) -> impl Iterator + use<'a> { + let case_sensitive = self.case_sensitive; + self.tokenizer.tokenize(text).into_iter().map(move |token| { + if case_sensitive { + element_hash(token.as_bytes()) + } else if token.is_ascii() { + buf.clear(); + buf.extend(token.bytes().map(|b| b.to_ascii_lowercase())); + element_hash(buf) + } else { + element_hash(token.to_lowercase().as_bytes()) + } + }) + } + /// Analyzes the given text into a list of tokens. pub fn analyze_text(&self, text: &str) -> Result> { let res = self @@ -162,6 +187,28 @@ impl Analyzer { mod tests { use super::*; + #[test] + fn test_analyze_text_hashes_matches_analyze_text() { + let text = "Hello, WORLD ship_Ship 清洁表面 ÄÖÜ straße İstanbul x"; + for (tokenizer, case_sensitive) in [ + (Box::new(EnglishTokenizer) as Box, false), + (Box::new(EnglishTokenizer), true), + (Box::new(ChineseTokenizer), false), + ] { + let analyzer = Analyzer::new(tokenizer, case_sensitive); + let expected = analyzer + .analyze_text(text) + .unwrap() + .iter() + .map(|t| element_hash(t)) + .collect::>(); + let hashes = analyzer + .analyze_text_hashes(text, &mut Vec::new()) + .collect::>(); + assert_eq!(expected, hashes); + } + } + #[test] fn test_english_tokenizer() { let tokenizer = EnglishTokenizer; @@ -215,6 +262,8 @@ mod tests { let text = "[2026/04/09/ 13:56:11.031]2026-04-09 13:56:11.031 - [ trace_id=340a6a44b0bd8e37bb7697ss7da61ff0 span_id=085ff5ttf1e0a23b trace_flags=01] - [http-nio-8081-exec-16] INFO c.h.p.xx.web.service.impl.CCCXForwardKKKServiceImpl.pushout(188) - 登录手机号18888888888的动态key：829889AC8 ship_ship ship__ship _ __ __IDENTIFIER__ _ship ship_ EOF"; let tokens = tokenizer.tokenize(text); + // Jieba 0.11 preserves ASCII compounds and emits searchable alphabetic parts. + // The underscore fallback also retains leading, trailing, and repeated underscores. assert_eq!( tokens, vec![ @@ -224,36 +273,34 @@ mod tests { "13", "56", "11.031", - "2026-04", - "09", + "2026-04-09", "13", "56", "11.031", "trace", - "_", "id", + "trace_id", "340a6a44b0bd8e37bb7697ss7da61ff0", "span", - "_", "id", + "span_id", "085ff5ttf1e0a23b", "trace", - "_", "flags", + "trace_flags", "01", "http", - "nio-8081", - "exec-16", + "nio", + "exec", + "http-nio-8081-exec-16", "INFO", - "c", - "h", - "p", "xx", "web", "service", "impl", "CCCXForwardKKKServiceImpl", "pushout", + "c.h.p.xx.web.service.impl.CCCXForwardKKKServiceImpl.pushout", "188", "登录", "手机", @@ -264,8 +311,8 @@ mod tests { "key", "829889AC8", "ship", - "_", "ship", + "ship_ship", "ship", "__", "ship", @@ -320,8 +367,8 @@ mod tests { "_", "id", "trace", - "_", "id", + "trace_id", "手机", "手机号", "_", diff --git a/src/index/src/inverted_index/create.rs b/src/index/src/inverted_index/create.rs index dade9705ac8..74da3612232 100644 --- a/src/index/src/inverted_index/create.rs +++ b/src/index/src/inverted_index/create.rs @@ -27,30 +27,24 @@ use crate::inverted_index::format::writer::InvertedIndexWriter; pub trait InvertedIndexCreator: Send { /// Adds a value to the named index. A `None` value represents an absence of data (null) /// - /// - `index_name`: Identifier for the index being built - /// - `value`: The data to be indexed, or `None` for a null entry - /// /// It should be equivalent to calling `push_with_name_n` with `n = 1` - async fn push_with_name( - &mut self, - index_name: &str, - value: Option>, - ) -> Result<()> { - self.push_with_name_n(index_name, value, 1).await + #[must_use = "a true result requires calling `spill` before pushing more"] + fn push_with_name(&mut self, index_name: &str, value: Option>) -> bool { + self.push_with_name_n(index_name, value, 1) } - /// Adds `n` identical values to the named index. `None` values represent absence of data (null) + /// Buffers `n` identical values for the named index. `None` values represent absence of + /// data (null). /// - /// - `index_name`: Identifier for the index being built - /// - `value`: The data to be indexed, or `None` for a null entry - /// - /// It should be equivalent to calling `push_with_name` `n` times - async fn push_with_name_n( - &mut self, - index_name: &str, - value: Option>, - n: usize, - ) -> Result<()>; + /// Returns true when buffered data exceeds the memory limit; the caller must then call + /// [`InvertedIndexCreator::spill`] before pushing more. Pushing is synchronous so the + /// per-row path does not allocate a future. + #[must_use = "a true result requires calling `spill` before pushing more"] + fn push_with_name_n(&mut self, index_name: &str, value: Option>, n: usize) + -> bool; + + /// Moves the buffers that asked for it to external storage. + async fn spill(&mut self) -> Result<()>; /// Finalizes the index creation process, ensuring all data is properly indexed and stored /// in the provided writer diff --git a/src/index/src/inverted_index/create/sort.rs b/src/index/src/inverted_index/create/sort.rs index bdc1ec21f62..fbf81adbaa4 100644 --- a/src/index/src/inverted_index/create/sort.rs +++ b/src/index/src/inverted_index/create/sort.rs @@ -42,15 +42,15 @@ pub struct SortOutput { /// Handles data sorting, supporting incremental input and retrieval of sorted output #[async_trait] pub trait Sorter: Send { - /// Inputs a non-null or null value into the sorter. - /// Should be equivalent to calling `push_n` with n = 1 - async fn push(&mut self, value: Option>) -> Result<()> { - self.push_n(value, 1).await - } + /// Buffers `n` identical non-null or null values in memory. + /// + /// Returns true when the buffer should be spilled with [`Sorter::spill`] before more + /// values are pushed. Kept synchronous so the per-row path does not allocate a future. + #[must_use = "a true result requires calling `spill` before pushing more"] + fn push_n(&mut self, value: Option>, n: usize) -> bool; - /// Pushing n identical non-null or null values into the sorter. - /// Should be equivalent to calling `push` n times - async fn push_n(&mut self, value: Option>, n: usize) -> Result<()>; + /// Moves the in-memory buffer to external storage. + async fn spill(&mut self) -> Result<()>; /// Completes the sorting process and returns the sorted data async fn output(&mut self) -> Result; diff --git a/src/index/src/inverted_index/create/sort/external_sort.rs b/src/index/src/inverted_index/create/sort/external_sort.rs index 12c49ca89de..cef07df9f8f 100644 --- a/src/index/src/inverted_index/create/sort/external_sort.rs +++ b/src/index/src/inverted_index/create/sort/external_sort.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{BTreeMap, VecDeque}; +use std::collections::{HashMap, VecDeque}; use std::mem; use std::num::NonZeroUsize; use std::ops::RangeInclusive; @@ -22,6 +22,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use async_trait::async_trait; use common_telemetry::{debug, error}; use futures::stream; +use roaring::RoaringBitmap; use snafu::ResultExt; use crate::bitmap::Bitmap; @@ -35,6 +36,28 @@ use crate::inverted_index::create::sort_create::SorterFactory; use crate::inverted_index::error::{IntermediateSnafu, Result}; use crate::{Bytes, BytesRef}; +/// Segments of one value. Rows arrive in order, so segment ids only grow. +struct Posting { + segments: RoaringBitmap, + last_segment: u32, +} + +/// Estimated cost of one segment id in a roaring array container. +const SEGMENT_SIZE: usize = size_of::(); + +impl Posting { + /// Adds `start..=end` and returns the estimated memory growth. + fn push(&mut self, start: u32, end: u32) -> usize { + if end <= self.last_segment { + return 0; + } + let start = start.max(self.last_segment + 1); + self.segments.insert_range(start..=end); + self.last_segment = end; + (end - start + 1) as usize * SEGMENT_SIZE + } +} + /// `ExternalSorter` manages the sorting of data using both in-memory structures and external files. /// It dumps data to external files when the in-memory buffer crosses a certain memory threshold. pub struct ExternalSorter { @@ -48,7 +71,7 @@ pub struct ExternalSorter { segment_null_bitmap: Bitmap, /// In-memory buffer to hold values and their corresponding bitmaps until memory threshold is exceeded - values_buffer: BTreeMap, + values_buffer: HashMap, /// Count of all rows ingested so far total_row_count: usize, @@ -78,11 +101,11 @@ pub struct ExternalSorter { #[async_trait] impl Sorter for ExternalSorter { - /// Pushes n identical values into the sorter, adding them to the in-memory buffer and dumping - /// the buffer to an external file if necessary - async fn push_n(&mut self, value: Option>, n: usize) -> Result<()> { + /// Pushes n identical values into the in-memory buffer; returns whether it should be + /// spilled. + fn push_n(&mut self, value: Option>, n: usize) -> bool { if n == 0 { - return Ok(()); + return false; } let segment_index_range = self.segment_index_range(n); @@ -90,13 +113,17 @@ impl Sorter for ExternalSorter { if let Some(value) = value { let memory_diff = self.push_not_null(value, segment_index_range); - self.may_dump_buffer(memory_diff).await + self.account_memory(memory_diff) } else { self.segment_null_bitmap.insert_range(segment_index_range); - Ok(()) + false } } + async fn spill(&mut self) -> Result<()> { + self.dump_buffer().await + } + /// Finalizes the sorting operation, merging data from both in-memory buffer and external files /// into a sorted stream async fn output(&mut self) -> Result { @@ -110,9 +137,7 @@ impl Sorter for ExternalSorter { let mut tree_nodes: VecDeque = VecDeque::with_capacity(readers.len() + 1); tree_nodes.push_back(Box::new(stream::iter( - mem::take(&mut self.values_buffer) - .into_iter() - .map(|(value, (bitmap, _))| Ok((value, bitmap))), + Self::sorted(mem::take(&mut self.values_buffer)).map(Ok), ))); for (_, reader) in readers { tree_nodes.push_back(IntermediateReader::new(reader).into_stream().await?); @@ -149,7 +174,7 @@ impl ExternalSorter { temp_file_provider, segment_null_bitmap: Bitmap::new_bitvec(), // bitvec is more efficient for many null values - values_buffer: BTreeMap::new(), + values_buffer: HashMap::default(), total_row_count: 0, segment_row_count, @@ -188,51 +213,65 @@ impl ExternalSorter { value: BytesRef<'_>, segment_index_range: RangeInclusive, ) -> usize { + let (start, end) = ( + *segment_index_range.start() as u32, + *segment_index_range.end() as u32, + ); match self.values_buffer.get_mut(value) { - Some((bitmap, mem_usage)) => { - bitmap.insert_range(segment_index_range); - let new_usage = bitmap.memory_usage() + value.len(); - let diff = new_usage - *mem_usage; - *mem_usage = new_usage; - - diff - } + Some(posting) => posting.push(start, end), None => { - let mut bitmap = Bitmap::new_roaring(); - bitmap.insert_range(segment_index_range); - - let mem_usage = bitmap.memory_usage() + value.len(); - self.values_buffer - .insert(value.to_vec(), (bitmap, mem_usage)); - - mem_usage + let mut segments = RoaringBitmap::new(); + segments.insert_range(start..=end); + let posting = Posting { + segments, + last_segment: end, + }; + self.values_buffer.insert(value.to_vec(), posting); + value.len() + (end - start + 1) as usize * SEGMENT_SIZE } } } - /// Checks if the in-memory buffer exceeds the threshold and offloads it to external storage if necessary - async fn may_dump_buffer(&mut self, memory_diff: usize) -> Result<()> { + /// Drains `values` sorted by value. + fn sorted( + values: HashMap, + ) -> impl Iterator { + let mut values = values.into_iter().collect::>(); + values.sort_unstable_by(|a, b| a.0.cmp(&b.0)); + values + .into_iter() + .map(|(value, posting)| (value, Bitmap::Roaring(posting.segments))) + } + + /// Records `memory_diff` and returns whether the buffer exceeds the thresholds and + /// should be offloaded to external storage. + fn account_memory(&mut self, memory_diff: usize) -> bool { self.current_memory_usage += memory_diff; let memory_usage = self.current_memory_usage; self.global_memory_usage .fetch_add(memory_diff, Ordering::Relaxed); - if self.global_memory_usage_sort_limit.is_none() { - return Ok(()); + let Some(limit) = self.global_memory_usage_sort_limit else { + return false; + }; + if self.global_memory_usage.load(Ordering::Relaxed) < limit { + return false; } - - if self.global_memory_usage.load(Ordering::Relaxed) - < self.global_memory_usage_sort_limit.unwrap() - { - return Ok(()); - } - if let Some(current_threshold) = self.current_memory_usage_threshold && memory_usage < current_threshold { + return false; + } + true + } + + async fn dump_buffer(&mut self) -> Result<()> { + // A second spill request before any new value would write an empty file with the + // same id and replace the first one. + if self.values_buffer.is_empty() { return Ok(()); } - + let memory_usage = self.current_memory_usage; let file_id = &format!("{:012}", self.total_row_count); let index_name = &self.index_name; let writer = self @@ -247,7 +286,7 @@ impl ExternalSorter { self.current_memory_usage = 0; let entries = values.len(); - IntermediateWriter::new(writer).write_all(values.into_iter().map(|(k, (b, _))| (k, b))).await.inspect(|_| + IntermediateWriter::new(writer).write_all(Self::sorted(values)).await.inspect(|_| debug!("Dumped {entries} entries ({memory_usage} bytes) to intermediate file {file_id} for index {index_name}") ).inspect_err(|e| error!(e; "Failed to dump {entries} entries to intermediate file {file_id} for index {index_name}") @@ -271,7 +310,7 @@ impl ExternalSorter { #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::collections::{BTreeMap, HashMap}; use std::iter; use std::sync::Mutex; @@ -329,7 +368,9 @@ mod tests { dictionary_values_and_sorted_result(row_count, segment_row_count); for (value, n) in dic_values { - sorter.push_n(value.as_deref(), n).await.unwrap(); + if sorter.push_n(value.as_deref(), n) { + sorter.spill().await.unwrap(); + } } sorted_result @@ -338,7 +379,9 @@ mod tests { shuffle_values_and_sorted_result(row_count, segment_row_count); for value in mock_values { - sorter.push(value.as_deref()).await.unwrap(); + if sorter.push_n(value.as_deref(), 1) { + sorter.spill().await.unwrap(); + } } sorted_result diff --git a/src/index/src/inverted_index/create/sort_create.rs b/src/index/src/inverted_index/create/sort_create.rs index cb22179a698..9218dd76ab6 100644 --- a/src/index/src/inverted_index/create/sort_create.rs +++ b/src/index/src/inverted_index/create/sort_create.rs @@ -42,6 +42,9 @@ pub struct SortIndexCreator { /// Number of rows in each segment, used to produce sorters segment_row_count: NonZeroUsize, + + /// Indexes whose sorter asked to spill its buffer. + pending_spills: Vec, } #[async_trait] @@ -50,22 +53,33 @@ impl InvertedIndexCreator for SortIndexCreator { /// /// If the index does not exist, a new index is created even if `n` is 0. /// Caller may leverage this behavior to create indexes with no data. - async fn push_with_name_n( + fn push_with_name_n( &mut self, index_name: &str, value: Option>, n: usize, - ) -> Result<()> { - match self.sorters.get_mut(index_name) { - Some(sorter) => sorter.push_n(value, n).await, + ) -> bool { + let sorter = match self.sorters.get_mut(index_name) { + Some(sorter) => sorter, None => { - let index_name = index_name.to_string(); - let mut sorter = (self.sorter_factory)(index_name.clone(), self.segment_row_count); - sorter.push_n(value, n).await?; - self.sorters.insert(index_name, sorter); - Ok(()) + let sorter = (self.sorter_factory)(index_name.to_string(), self.segment_row_count); + self.sorters.entry(index_name.to_string()).or_insert(sorter) + } + }; + if sorter.push_n(value, n) { + self.pending_spills.push(index_name.to_string()); + return true; + } + false + } + + async fn spill(&mut self) -> Result<()> { + for index_name in std::mem::take(&mut self.pending_spills) { + if let Some(sorter) = self.sorters.get_mut(&index_name) { + sorter.spill().await?; } } + Ok(()) } /// Finalizes the sorting for all indexes and writes them using the inverted index writer @@ -110,6 +124,7 @@ impl SortIndexCreator { sorter_factory, sorters: HashMap::new(), segment_row_count, + pending_spills: Vec::new(), } } } @@ -117,6 +132,7 @@ impl SortIndexCreator { #[cfg(test)] mod tests { use std::collections::BTreeMap; + use std::sync::{Arc, Mutex}; use common_base::BitVec; use futures::{StreamExt, stream}; @@ -140,10 +156,7 @@ mod tests { for (index_name, values) in index_values { for value in values { - creator - .push_with_name(index_name, Some(value)) - .await - .unwrap(); + assert!(!creator.push_with_name(index_name, Some(value))); } } @@ -189,10 +202,7 @@ mod tests { for (index_name, values) in index_values { for value in values { - creator - .push_with_name(index_name, Some(value)) - .await - .unwrap(); + assert!(!creator.push_with_name(index_name, Some(value))); } } @@ -221,9 +231,9 @@ mod tests { let mut creator = SortIndexCreator::new(NaiveSorter::factory(), NonZeroUsize::new(1).unwrap()); - creator.push_with_name_n("a", None, 0).await.unwrap(); - creator.push_with_name_n("b", None, 0).await.unwrap(); - creator.push_with_name_n("c", None, 0).await.unwrap(); + assert!(!creator.push_with_name_n("a", None, 0)); + assert!(!creator.push_with_name_n("b", None, 0)); + assert!(!creator.push_with_name_n("c", None, 0)); let mut mock_writer = MockInvertedIndexWriter::new(); mock_writer @@ -250,6 +260,51 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn test_sort_index_creator_spills_requested_sorters() { + let spilled = Arc::new(Mutex::new(Vec::new())); + let factory: SorterFactory = { + let spilled = spilled.clone(); + Box::new(move |index_name, _| { + Box::new(SpillRecordingSorter { + index_name, + spilled: spilled.clone(), + }) + }) + }; + let mut creator = SortIndexCreator::new(factory, NonZeroUsize::new(1).unwrap()); + + assert!(creator.push_with_name("a", Some(b"1"))); + assert!(!creator.push_with_name("b", Some(b"1"))); + creator.spill().await.unwrap(); + assert_eq!(*spilled.lock().unwrap(), vec!["a"]); + + creator.spill().await.unwrap(); + assert_eq!(*spilled.lock().unwrap(), vec!["a"]); + } + + /// Requests a spill on every push to index `a` and records the spills it receives. + struct SpillRecordingSorter { + index_name: String, + spilled: Arc>>, + } + + #[async_trait] + impl Sorter for SpillRecordingSorter { + fn push_n(&mut self, _value: Option>, _n: usize) -> bool { + self.index_name == "a" + } + + async fn spill(&mut self) -> Result<()> { + self.spilled.lock().unwrap().push(self.index_name.clone()); + Ok(()) + } + + async fn output(&mut self) -> Result { + unreachable!() + } + } + fn set_bit(bit_vec: &mut BitVec, index: usize) { if index >= bit_vec.len() { bit_vec.resize(index + 1, false); @@ -277,20 +332,18 @@ mod tests { #[async_trait] impl Sorter for NaiveSorter { - async fn push(&mut self, value: Option>) -> Result<()> { - let segment_index = self.total_row_count / self.segment_row_count; - self.total_row_count += 1; + fn push_n(&mut self, value: Option>, n: usize) -> bool { + for _ in 0..n { + let segment_index = self.total_row_count / self.segment_row_count; + self.total_row_count += 1; - let bitmap = self.values.entry(value.map(Into::into)).or_default(); - set_bit(bitmap, segment_index); - - Ok(()) + let bitmap = self.values.entry(value.map(Into::into)).or_default(); + set_bit(bitmap, segment_index); + } + false } - async fn push_n(&mut self, value: Option>, n: usize) -> Result<()> { - for _ in 0..n { - self.push(value).await?; - } + async fn spill(&mut self) -> Result<()> { Ok(()) } diff --git a/src/index/src/inverted_index/search.rs b/src/index/src/inverted_index/search.rs index 8e28440c7e2..581135c2d4e 100644 --- a/src/index/src/inverted_index/search.rs +++ b/src/index/src/inverted_index/search.rs @@ -16,3 +16,13 @@ pub mod fst_apply; pub mod fst_values_mapper; pub mod index_apply; pub mod predicate; + +/// Partitions `slice` in place so that elements matching `pred` come first, +/// preserving the original relative order within each group. Returns the +/// number of matching elements. +/// +/// Stable replacement for the unstable `Iterator::partition_in_place`. +pub(crate) fn partition_in_place(slice: &mut [T], mut pred: impl FnMut(&T) -> bool) -> usize { + slice.sort_by_key(|x| !pred(x)); + slice.iter().filter(|x| pred(x)).count() +} diff --git a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs b/src/index/src/inverted_index/search/fst_apply/keys_apply.rs index d565749452d..b68f58cdfd3 100644 --- a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs +++ b/src/index/src/inverted_index/search/fst_apply/keys_apply.rs @@ -68,16 +68,16 @@ impl KeysFstApplier { } fn split_at_in_lists(predicates: &mut [Predicate]) -> (&mut [Predicate], &mut [Predicate]) { - let in_list_index = predicates - .iter_mut() - .partition_in_place(|p| matches!(p, Predicate::InList(_))); + let in_list_index = crate::inverted_index::search::partition_in_place(predicates, |p| { + matches!(p, Predicate::InList(_)) + }); predicates.split_at_mut(in_list_index) } fn split_at_ranges(predicates: &mut [Predicate]) -> (&mut [Predicate], &mut [Predicate]) { - let range_index = predicates - .iter_mut() - .partition_in_place(|p| matches!(p, Predicate::Range(_))); + let range_index = crate::inverted_index::search::partition_in_place(predicates, |p| { + matches!(p, Predicate::Range(_)) + }); predicates.split_at_mut(range_index) } diff --git a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs index eaea8cfb95c..5fcc962ca7d 100644 --- a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs +++ b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs @@ -127,9 +127,10 @@ impl PredicatesIndexApplier { let mut fst_appliers = Vec::with_capacity(predicates.len()); // InList predicates are applied first to benefit from higher selectivity. - let in_list_index = predicates - .iter_mut() - .partition_in_place(|(_, ps)| ps.iter().any(|p| matches!(p, Predicate::InList(_)))); + let in_list_index = + crate::inverted_index::search::partition_in_place(&mut predicates, |(_, ps)| { + ps.iter().any(|p| matches!(p, Predicate::InList(_))) + }); let mut iter = predicates.into_iter(); for _ in 0..in_list_index { let (column_name, predicates) = iter.next().unwrap(); diff --git a/src/index/src/lib.rs b/src/index/src/lib.rs index c3ae35e9072..571866ade90 100644 --- a/src/index/src/lib.rs +++ b/src/index/src/lib.rs @@ -12,9 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(iter_partition_in_place)] -#![feature(hash_set_entry)] - pub mod bitmap; pub mod bloom_filter; pub mod error; @@ -22,8 +19,6 @@ pub mod external_provider; pub mod fulltext_index; pub mod inverted_index; pub mod target; -#[cfg(feature = "vector_index")] -pub mod vector; pub type Bytes = Vec; pub type BytesRef<'a> = &'a [u8]; diff --git a/src/index/src/vector.rs b/src/index/src/vector.rs deleted file mode 100644 index 9c7cfded330..00000000000 --- a/src/index/src/vector.rs +++ /dev/null @@ -1,179 +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. - -//! Vector index implementation for KNN search. -//! -//! This module provides the core algorithms for vector index creation, reading, -//! and search operations. The architecture follows the same pattern as other indexes: -//! -//! - `error.rs`: Error types for vector index operations -//! - `format.rs`: Proto conversion helpers for blob format -//! - `util.rs`: Byte utilities for vector data -//! - `engine/`: Pluggable vector index engines (e.g., USearch HNSW) -//! - `create.rs`: `VectorIndexCreator` trait and implementations for building indices -//! - `apply.rs`: `VectorIndexApplier` trait and implementations for searching indices - -pub mod apply; -pub mod create; -pub mod engine; -pub mod error; -pub mod format; -pub mod util; - -pub use apply::{HnswVectorIndexApplier, IndexApplier, VectorSearchOutput}; -pub use create::{HnswVectorIndexCreator, IndexCreator}; -pub use datatypes::schema::{VectorDistanceMetric, VectorIndexOptions}; -use nalgebra::DVectorView; -pub use usearch::MetricKind; - -/// Converts a VectorDistanceMetric to a USearch MetricKind. -pub fn distance_metric_to_usearch(metric: VectorDistanceMetric) -> MetricKind { - match metric { - VectorDistanceMetric::L2sq => MetricKind::L2sq, - VectorDistanceMetric::Cosine => MetricKind::Cos, - VectorDistanceMetric::InnerProduct => MetricKind::IP, - } -} - -/// Computes distance between two vectors using the specified metric. -/// -/// Uses SIMD-optimized implementations via nalgebra. -/// -/// **Note:** The caller must ensure that the two vectors have the same length -/// and are non-empty. Empty vectors return 0.0 for all metrics. -pub fn compute_distance(v1: &[f32], v2: &[f32], metric: VectorDistanceMetric) -> f32 { - // Empty vectors are degenerate; return 0.0 uniformly across all metrics. - if v1.is_empty() || v2.is_empty() { - return 0.0; - } - - match metric { - VectorDistanceMetric::L2sq => l2sq(v1, v2), - VectorDistanceMetric::Cosine => cosine(v1, v2), - VectorDistanceMetric::InnerProduct => -dot(v1, v2), - } -} - -/// Calculates the squared L2 distance between two vectors. -fn l2sq(lhs: &[f32], rhs: &[f32]) -> f32 { - let lhs = DVectorView::from_slice(lhs, lhs.len()); - let rhs = DVectorView::from_slice(rhs, rhs.len()); - (lhs - rhs).norm_squared() -} - -/// Calculates the cosine distance between two vectors. -/// -/// Returns a value in `[0.0, 2.0]` where 0.0 means identical direction and 2.0 means -/// opposite direction. For degenerate cases (zero or near-zero magnitude vectors), -/// returns 1.0 (maximum uncertainty) to avoid NaN and ensure safe index operations. -fn cosine(lhs: &[f32], rhs: &[f32]) -> f32 { - let lhs_vec = DVectorView::from_slice(lhs, lhs.len()); - let rhs_vec = DVectorView::from_slice(rhs, rhs.len()); - - let dot_product = lhs_vec.dot(&rhs_vec); - let lhs_norm = lhs_vec.norm(); - let rhs_norm = rhs_vec.norm(); - - // Zero-magnitude vectors have undefined direction; return max distance as safe fallback. - if dot_product.abs() < f32::EPSILON - || lhs_norm.abs() < f32::EPSILON - || rhs_norm.abs() < f32::EPSILON - { - return 1.0; - } - - let cos_similar = dot_product / (lhs_norm * rhs_norm); - let res = 1.0 - cos_similar; - // Clamp near-zero results to exactly 0.0 to avoid floating-point artifacts. - if res.abs() < f32::EPSILON { 0.0 } else { res } -} - -/// Calculates the dot product between two vectors. -fn dot(lhs: &[f32], rhs: &[f32]) -> f32 { - let lhs = DVectorView::from_slice(lhs, lhs.len()); - let rhs = DVectorView::from_slice(rhs, rhs.len()); - lhs.dot(&rhs) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_distance_metric_to_usearch() { - assert_eq!( - distance_metric_to_usearch(VectorDistanceMetric::L2sq), - MetricKind::L2sq - ); - assert_eq!( - distance_metric_to_usearch(VectorDistanceMetric::Cosine), - MetricKind::Cos - ); - assert_eq!( - distance_metric_to_usearch(VectorDistanceMetric::InnerProduct), - MetricKind::IP - ); - } - - #[test] - fn test_vector_index_options_default() { - let options = VectorIndexOptions::default(); - assert_eq!(options.metric, VectorDistanceMetric::L2sq); - assert_eq!(options.connectivity, 16); - assert_eq!(options.expansion_add, 128); - assert_eq!(options.expansion_search, 64); - } - - #[test] - fn test_compute_distance_l2sq() { - let v1 = vec![1.0, 2.0, 3.0]; - let v2 = vec![4.0, 5.0, 6.0]; - // L2sq = (4-1)^2 + (5-2)^2 + (6-3)^2 = 9 + 9 + 9 = 27 - let dist = compute_distance(&v1, &v2, VectorDistanceMetric::L2sq); - assert!((dist - 27.0).abs() < 1e-6); - } - - #[test] - fn test_compute_distance_cosine() { - let v1 = vec![1.0, 0.0, 0.0]; - let v2 = vec![0.0, 1.0, 0.0]; - // Orthogonal vectors have cosine similarity of 0, distance of 1 - let dist = compute_distance(&v1, &v2, VectorDistanceMetric::Cosine); - assert!((dist - 1.0).abs() < 1e-6); - } - - #[test] - fn test_compute_distance_inner_product() { - let v1 = vec![1.0, 2.0, 3.0]; - let v2 = vec![4.0, 5.0, 6.0]; - // Inner product = 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32 - // Distance is negated: -32 - let dist = compute_distance(&v1, &v2, VectorDistanceMetric::InnerProduct); - assert!((dist - (-32.0)).abs() < 1e-6); - } - - #[test] - fn test_compute_distance_empty_vectors() { - // Empty vectors should return 0.0 uniformly for all metrics - assert_eq!(compute_distance(&[], &[], VectorDistanceMetric::L2sq), 0.0); - assert_eq!( - compute_distance(&[], &[], VectorDistanceMetric::Cosine), - 0.0 - ); - assert_eq!( - compute_distance(&[], &[], VectorDistanceMetric::InnerProduct), - 0.0 - ); - } -} diff --git a/src/index/src/vector/apply.rs b/src/index/src/vector/apply.rs deleted file mode 100644 index 82db03b9efd..00000000000 --- a/src/index/src/vector/apply.rs +++ /dev/null @@ -1,428 +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. - -//! Vector index application trait and implementation. - -use greptime_proto::v1::index::VectorIndexMeta; -use prost::Message; -use roaring::RoaringBitmap; -use snafu::ResultExt; -use store_api::storage::VectorIndexEngine; - -use crate::vector::VectorDistanceMetric; -use crate::vector::engine::{self, VectorIndexConfig}; -use crate::vector::error::{ - BlobTooSmallSnafu, BlobTruncatedSnafu, DecodeProtoSnafu, DeserializeBitmapSnafu, EngineSnafu, - InvalidBlobSnafu, KeyMappingSnafu, Result, UnknownDistanceMetricSnafu, UnknownEngineTypeSnafu, -}; -use crate::vector::format::{META_SIZE_LEN, distance_metric_from_proto, engine_type_from_proto}; - -/// Output of a vector index search. -#[derive(Debug, Clone)] -pub struct VectorSearchOutput { - /// Row offsets in the original data. - pub row_offsets: Vec, - /// Distances to query vector. - pub distances: Vec, -} - -/// `IndexApplier` is for searching a vector index. -pub trait IndexApplier: Send + Sync { - /// Searches for k nearest neighbors. - fn search(&self, query: &[f32], k: usize) -> Result; - - /// Returns the memory usage. - fn memory_usage(&self) -> usize; - - /// Returns the dimension of vectors in this index. - fn dimensions(&self) -> u32; - - /// Returns the distance metric used for this index. - fn metric(&self) -> VectorDistanceMetric; - - /// Returns the total number of rows (including NULLs). - fn total_rows(&self) -> u64; - - /// Returns the number of indexed rows (excluding NULLs). - fn indexed_rows(&self) -> u64; -} - -/// HNSW-based vector index applier. -pub struct HnswVectorIndexApplier { - /// The loaded vector index engine. - engine: Box, - /// Bitmap tracking which row offsets have NULL vectors. - null_bitmap: RoaringBitmap, - /// The dimension of vectors in this index. - dimensions: u32, - /// The distance metric used for this index. - metric: VectorDistanceMetric, - /// Total number of rows (including NULLs). - total_rows: u64, - /// Number of indexed rows (excluding NULLs). - indexed_rows: u64, -} - -impl HnswVectorIndexApplier { - /// Creates a new HNSW vector index applier from blob data. - /// - /// Blob layout: - /// ```text - /// +----------------------+ - /// | Null Bitmap | <- RoaringBitmap serialized - /// +----------------------+ - /// | Index Data | <- engine-specific serialized - /// +----------------------+ - /// | Meta (protobuf) | <- VectorIndexMeta encoded - /// +----------------------+ - /// | Meta Size (4B) | <- u32 LE - /// +----------------------+ - /// ``` - pub fn from_blob(data: &[u8]) -> Result { - if data.len() < META_SIZE_LEN { - return BlobTooSmallSnafu { - min_size: META_SIZE_LEN, - actual_size: data.len(), - } - .fail(); - } - - // Read meta size from the end - let blob_len = data.len(); - let meta_size = - u32::from_le_bytes(data[blob_len - META_SIZE_LEN..].try_into().unwrap()) as usize; - - if blob_len < META_SIZE_LEN + meta_size { - return BlobTruncatedSnafu { - reason: "blob truncated while reading meta", - } - .fail(); - } - - // Read and decode meta - let meta_start = blob_len - META_SIZE_LEN - meta_size; - let meta = VectorIndexMeta::decode(&data[meta_start..blob_len - META_SIZE_LEN]) - .context(DecodeProtoSnafu)?; - - let null_bitmap_size = meta.null_bitmap_size as usize; - let index_size = meta.index_size as usize; - - if null_bitmap_size + index_size > meta_start { - return InvalidBlobSnafu { - reason: "invalid size fields", - } - .fail(); - } - - // Read null bitmap - let null_bitmap = RoaringBitmap::deserialize_from(&data[..null_bitmap_size]) - .context(DeserializeBitmapSnafu)?; - - // Read index data - let index_start = null_bitmap_size; - let index_end = index_start + index_size; - let index_bytes = &data[index_start..index_end]; - - let engine_proto = meta.engine.try_into().map_err(|_| { - UnknownEngineTypeSnafu { - engine_type: meta.engine, - } - .build() - })?; - let engine_type = engine_type_from_proto(engine_proto).ok_or_else(|| { - UnknownEngineTypeSnafu { - engine_type: meta.engine, - } - .build() - })?; - - let metric_proto = meta.metric.try_into().map_err(|_| { - UnknownDistanceMetricSnafu { - metric: meta.metric, - } - .build() - })?; - let distance_metric = distance_metric_from_proto(metric_proto).ok_or_else(|| { - UnknownDistanceMetricSnafu { - metric: meta.metric, - } - .build() - })?; - - let config = VectorIndexConfig { - engine: engine_type, - dim: meta.dim as usize, - distance_metric, - connectivity: meta.connectivity as usize, - expansion_add: meta.expansion_add as usize, - expansion_search: meta.expansion_search as usize, - }; - - let engine = engine::load_engine(engine_type, &config, index_bytes)?; - - let stats = meta.stats.ok_or_else(|| { - InvalidBlobSnafu { - reason: "missing stats in vector index metadata", - } - .build() - })?; - - Ok(Self { - engine, - null_bitmap, - dimensions: meta.dim, - metric: distance_metric, - total_rows: stats.total_row_count, - indexed_rows: stats.indexed_row_count, - }) - } - - /// Maps HNSW keys from search results to row offsets in the original data. - fn map_keys_to_row_offsets(&self, keys: Vec) -> Result> { - if self.total_rows == 0 { - return Ok(Vec::new()); - } - let total_rows_u32 = u32::try_from(self.total_rows).map_err(|_| { - KeyMappingSnafu { - reason: format!("Total rows {} exceeds u32::MAX", self.total_rows), - } - .build() - })?; - - let mut row_offsets = Vec::with_capacity(keys.len()); - for key in keys { - let offset = self.hnsw_key_to_row_offset(total_rows_u32, key)?; - row_offsets.push(offset as u64); - } - Ok(row_offsets) - } - - /// Converts a single HNSW key to its corresponding row offset. - fn hnsw_key_to_row_offset(&self, total_rows: u32, key: u64) -> Result { - if total_rows == 0 { - return KeyMappingSnafu { - reason: "Total rows is zero", - } - .fail(); - } - if key >= self.indexed_rows { - return KeyMappingSnafu { - reason: format!( - "HNSW key {} exceeds indexed rows {}", - key, self.indexed_rows - ), - } - .fail(); - } - - // Fast path: no NULLs means key == row offset - if self.null_bitmap.is_empty() { - return Ok(key as u32); - } - - // Binary search to find the row offset - let mut left: u32 = 0; - let mut right: u32 = total_rows - 1; - while left <= right { - let mid = left + (right - left) / 2; - let nulls_before = self.null_bitmap.rank(mid); - let non_nulls = (mid as u64 + 1).saturating_sub(nulls_before); - if non_nulls > key { - if mid == 0 { - break; - } - right = mid - 1; - } else { - left = mid + 1; - } - } - - if left >= total_rows { - return KeyMappingSnafu { - reason: "Failed to map HNSW key to row offset", - } - .fail(); - } - - Ok(left) - } -} - -impl IndexApplier for HnswVectorIndexApplier { - fn search(&self, query: &[f32], k: usize) -> Result { - if self.indexed_rows == 0 || k == 0 { - return Ok(VectorSearchOutput { - row_offsets: Vec::new(), - distances: Vec::new(), - }); - } - - let actual_k = k.min(self.indexed_rows as usize); - let matches = self.engine.search(query, actual_k).map_err(|e| { - EngineSnafu { - reason: e.to_string(), - } - .build() - })?; - - let row_offsets = self.map_keys_to_row_offsets(matches.keys)?; - - Ok(VectorSearchOutput { - row_offsets, - distances: matches.distances, - }) - } - - fn memory_usage(&self) -> usize { - self.engine.memory_usage() - + self.null_bitmap.serialized_size() - + std::mem::size_of::() - } - - fn dimensions(&self) -> u32 { - self.dimensions - } - - fn metric(&self) -> VectorDistanceMetric { - self.metric - } - - fn total_rows(&self) -> u64 { - self.total_rows - } - - fn indexed_rows(&self) -> u64 { - self.indexed_rows - } -} - -#[cfg(test)] -mod tests { - use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; - use prost::Message; - use roaring::RoaringBitmap; - use store_api::storage::VectorIndexEngineType; - - use super::*; - use crate::vector::engine::VectorIndexConfig; - use crate::vector::format::{distance_metric_to_proto, engine_type_to_proto}; - - fn build_test_blob( - config: &VectorIndexConfig, - vectors: Vec<(u64, Vec)>, - null_bitmap: &RoaringBitmap, - total_rows: u64, - indexed_rows: u64, - ) -> Vec { - let mut engine = engine::create_engine(config.engine, config).unwrap(); - for (key, vector) in vectors { - engine.add(key, &vector).unwrap(); - } - let index_size = engine.serialized_length(); - let mut index_bytes = vec![0u8; index_size]; - engine.save_to_buffer(&mut index_bytes).unwrap(); - - let mut null_bitmap_bytes = Vec::new(); - null_bitmap.serialize_into(&mut null_bitmap_bytes).unwrap(); - - let null_count = total_rows - indexed_rows; - let meta = VectorIndexMeta { - engine: engine_type_to_proto(config.engine).into(), - dim: config.dim as u32, - metric: distance_metric_to_proto(config.distance_metric).into(), - connectivity: config.connectivity as u32, - expansion_add: config.expansion_add as u32, - expansion_search: config.expansion_search as u32, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_bytes.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: total_rows, - indexed_row_count: indexed_rows, - null_count, - }), - }; - let meta_bytes = meta.encode_to_vec(); - let meta_size = meta_bytes.len() as u32; - - let mut blob = Vec::new(); - blob.extend_from_slice(&null_bitmap_bytes); - blob.extend_from_slice(&index_bytes); - blob.extend_from_slice(&meta_bytes); - blob.extend_from_slice(&meta_size.to_le_bytes()); - blob - } - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 2, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[test] - fn test_hnsw_vector_index_applier_search() { - let config = test_config(); - let mut null_bitmap = RoaringBitmap::new(); - null_bitmap.insert(1); - - let blob = build_test_blob( - &config, - vec![(0, vec![1.0, 0.0]), (1, vec![0.0, 1.0])], - &null_bitmap, - 3, - 2, - ); - - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - assert_eq!(applier.dimensions(), 2); - assert_eq!(applier.metric(), VectorDistanceMetric::L2sq); - assert_eq!(applier.total_rows(), 3); - assert_eq!(applier.indexed_rows(), 2); - - let output = applier.search(&[1.0, 0.0], 2).unwrap(); - assert_eq!(output.row_offsets.len(), 2); - // First result should be row 0 (closest to [1.0, 0.0]) - assert_eq!(output.row_offsets[0], 0); - // Second result should be row 2 (maps from HNSW key 1) - assert_eq!(output.row_offsets[1], 2); - } - - #[test] - fn test_hnsw_vector_index_applier_empty() { - let config = test_config(); - let null_bitmap = RoaringBitmap::new(); - - let blob = build_test_blob(&config, vec![], &null_bitmap, 0, 0); - - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - let output = applier.search(&[1.0, 0.0], 1).unwrap(); - assert!(output.row_offsets.is_empty()); - } - - #[test] - fn test_hnsw_vector_index_applier_k_zero() { - let config = test_config(); - let null_bitmap = RoaringBitmap::new(); - - let blob = build_test_blob(&config, vec![(0, vec![1.0, 0.0])], &null_bitmap, 1, 1); - - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - let output = applier.search(&[1.0, 0.0], 0).unwrap(); - assert!(output.row_offsets.is_empty()); - } -} diff --git a/src/index/src/vector/create.rs b/src/index/src/vector/create.rs deleted file mode 100644 index 006bd786908..00000000000 --- a/src/index/src/vector/create.rs +++ /dev/null @@ -1,345 +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. - -//! Vector index creation trait and implementation. - -use async_trait::async_trait; -use common_error::ext::BoxedError; -use futures::io::Cursor; -use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; -use prost::Message; -use puffin::puffin_manager::{PuffinWriter, PutOptions}; -use roaring::RoaringBitmap; -use snafu::ResultExt; -use store_api::storage::{VectorIndexEngine, VectorIndexEngineType}; - -use crate::vector::engine::{self, VectorIndexConfig}; -use crate::vector::error::{ - EngineSnafu, ExternalSnafu, FinishSnafu, Result, RowCountOverflowSnafu, SerializeBitmapSnafu, -}; -use crate::vector::format::{META_SIZE_LEN, distance_metric_to_proto, engine_type_to_proto}; - -/// `IndexCreator` is for creating a vector index. -#[async_trait] -pub trait IndexCreator: Send { - /// Pushes a vector to the index. - fn push_vector(&mut self, vector: &[f32]) -> Result<()>; - - /// Records a NULL vector. - fn push_null(&mut self) -> Result<()>; - - /// Records multiple NULL vectors. - fn push_nulls(&mut self, n: usize) -> Result<()>; - - /// Finalizes the index and writes to puffin. - /// Returns bytes written. - async fn finish( - &mut self, - puffin_writer: &mut (impl PuffinWriter + Send), - blob_key: &str, - put_options: PutOptions, - ) -> Result; - - /// Aborts the creation. - async fn abort(&mut self) -> Result<()>; - - /// Returns the memory usage. - fn memory_usage(&self) -> usize; - - /// Returns the number of indexed vectors. - fn size(&self) -> usize; -} - -/// HNSW-based vector index creator. -pub struct HnswVectorIndexCreator { - /// The vector index engine (e.g., USearch HNSW). - engine: Box, - /// Configuration for this index. - config: VectorIndexConfig, - /// Bitmap tracking which row offsets have NULL vectors. - null_bitmap: RoaringBitmap, - /// Current row offset (including NULLs). - current_row_offset: u64, - /// Next HNSW key to assign (only for non-NULL vectors). - next_hnsw_key: u64, - /// Memory usage estimation. - memory_usage_bytes: usize, - /// Whether the index has been finished. - finished: bool, -} - -impl HnswVectorIndexCreator { - /// Creates a new HNSW vector index creator. - pub fn new(config: VectorIndexConfig) -> Result { - let engine_instance = engine::create_engine(config.engine, &config)?; - - Ok(Self { - engine: engine_instance, - config, - null_bitmap: RoaringBitmap::new(), - current_row_offset: 0, - next_hnsw_key: 0, - memory_usage_bytes: 0, - finished: false, - }) - } - - /// Reserves capacity for the expected number of vectors. - #[allow(dead_code)] - pub fn reserve(&mut self, capacity: usize) -> Result<()> { - self.engine.reserve(capacity).map_err(|e| { - EngineSnafu { - reason: format!("Failed to reserve capacity: {}", e), - } - .build() - }) - } - - /// Returns the engine type. - pub fn engine_type(&self) -> VectorIndexEngineType { - self.config.engine - } - - /// Returns the configuration. - pub fn config(&self) -> &VectorIndexConfig { - &self.config - } - - /// Returns the null bitmap. - pub fn null_bitmap(&self) -> &RoaringBitmap { - &self.null_bitmap - } - - /// Returns the current row offset (total rows including NULLs). - pub fn current_row_offset(&self) -> u64 { - self.current_row_offset - } - - /// Returns the next HNSW key (total indexed rows excluding NULLs). - pub fn next_hnsw_key(&self) -> u64 { - self.next_hnsw_key - } - - /// Builds the blob data for the vector index. - fn build_blob(&self) -> Result> { - // Serialize null bitmap - let mut null_bitmap_bytes = Vec::new(); - self.null_bitmap - .serialize_into(&mut null_bitmap_bytes) - .context(SerializeBitmapSnafu)?; - - // Serialize the vector index - let index_size = self.engine.serialized_length(); - let mut index_bytes = vec![0u8; index_size]; - self.engine.save_to_buffer(&mut index_bytes).map_err(|e| { - EngineSnafu { - reason: format!("Failed to serialize index: {}", e), - } - .build() - })?; - - // Build VectorIndexMeta - let total_rows = self.current_row_offset; - let indexed_rows = self.next_hnsw_key; - let null_count = total_rows - indexed_rows; - - let meta = VectorIndexMeta { - engine: engine_type_to_proto(self.config.engine).into(), - dim: self.config.dim as u32, - metric: distance_metric_to_proto(self.config.distance_metric).into(), - connectivity: self.config.connectivity as u32, - expansion_add: self.config.expansion_add as u32, - expansion_search: self.config.expansion_search as u32, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_bytes.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: total_rows, - indexed_row_count: indexed_rows, - null_count, - }), - }; - let meta_bytes = meta.encode_to_vec(); - let meta_size = meta_bytes.len() as u32; - - // Build blob: null_bitmap | index_data | meta | meta_size - let total_size = - null_bitmap_bytes.len() + index_bytes.len() + meta_bytes.len() + META_SIZE_LEN; - let mut blob_data = Vec::with_capacity(total_size); - blob_data.extend_from_slice(&null_bitmap_bytes); - blob_data.extend_from_slice(&index_bytes); - blob_data.extend_from_slice(&meta_bytes); - blob_data.extend_from_slice(&meta_size.to_le_bytes()); - - Ok(blob_data) - } -} - -#[async_trait] -impl IndexCreator for HnswVectorIndexCreator { - fn push_vector(&mut self, vector: &[f32]) -> Result<()> { - if self.current_row_offset >= u32::MAX as u64 { - return RowCountOverflowSnafu { - current: self.current_row_offset, - increment: 1u64, - } - .fail(); - } - - let key = self.next_hnsw_key; - self.engine.add(key, vector).map_err(|e| { - EngineSnafu { - reason: e.to_string(), - } - .build() - })?; - self.next_hnsw_key += 1; - self.current_row_offset += 1; - self.memory_usage_bytes = self.engine.memory_usage(); - Ok(()) - } - - fn push_null(&mut self) -> Result<()> { - if self.current_row_offset >= u32::MAX as u64 { - return RowCountOverflowSnafu { - current: self.current_row_offset, - increment: 1u64, - } - .fail(); - } - - self.null_bitmap.insert(self.current_row_offset as u32); - self.current_row_offset += 1; - Ok(()) - } - - fn push_nulls(&mut self, n: usize) -> Result<()> { - let increment = n as u64; - if self.current_row_offset.saturating_add(increment) > u32::MAX as u64 { - return RowCountOverflowSnafu { - current: self.current_row_offset, - increment, - } - .fail(); - } - - let start = self.current_row_offset as u32; - let end = (self.current_row_offset + increment) as u32; - self.null_bitmap.insert_range(start..end); - self.current_row_offset += increment; - Ok(()) - } - - async fn finish( - &mut self, - puffin_writer: &mut (impl PuffinWriter + Send), - blob_key: &str, - put_options: PutOptions, - ) -> Result { - if self.finished { - return FinishSnafu { - reason: "Index already finished", - } - .fail(); - } - self.finished = true; - - let blob_data = self.build_blob()?; - - puffin_writer - .put_blob( - blob_key, - Cursor::new(blob_data), - put_options, - Default::default(), - ) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - - async fn abort(&mut self) -> Result<()> { - self.finished = true; - Ok(()) - } - - fn memory_usage(&self) -> usize { - self.memory_usage_bytes + self.null_bitmap.serialized_size() - } - - fn size(&self) -> usize { - self.engine.size() - } -} - -#[cfg(test)] -mod tests { - use store_api::storage::VectorIndexEngineType; - - use super::*; - use crate::vector::VectorDistanceMetric; - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 4, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[test] - fn test_hnsw_vector_index_creator() { - let config = test_config(); - let mut creator = HnswVectorIndexCreator::new(config).unwrap(); - - creator.reserve(10).unwrap(); - - // Add some vectors - let v1 = vec![1.0f32, 0.0, 0.0, 0.0]; - let v2 = vec![0.0f32, 1.0, 0.0, 0.0]; - - creator.push_vector(&v1).unwrap(); - creator.push_null().unwrap(); - creator.push_vector(&v2).unwrap(); - - assert_eq!(creator.size(), 2); // 2 vectors (excluding NULL) - assert_eq!(creator.current_row_offset(), 3); // 3 rows total - assert!(creator.null_bitmap().contains(1)); // Row 1 is NULL - } - - #[test] - fn test_hnsw_vector_index_creator_build_blob() { - let config = test_config(); - let mut creator = HnswVectorIndexCreator::new(config).unwrap(); - - creator.reserve(10).unwrap(); - - // Add vectors - creator.push_vector(&[1.0, 0.0, 0.0, 0.0]).unwrap(); - creator.push_null().unwrap(); - creator.push_vector(&[0.0, 1.0, 0.0, 0.0]).unwrap(); - - let blob = creator.build_blob().unwrap(); - assert!(!blob.is_empty()); - - // Verify blob can be parsed using the applier - use crate::vector::apply::{HnswVectorIndexApplier, IndexApplier}; - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - assert_eq!(applier.dimensions(), 4); - assert_eq!(applier.total_rows(), 3); - assert_eq!(applier.indexed_rows(), 2); - } -} diff --git a/src/index/src/vector/engine.rs b/src/index/src/vector/engine.rs deleted file mode 100644 index b99ba89b96d..00000000000 --- a/src/index/src/vector/engine.rs +++ /dev/null @@ -1,63 +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. - -//! Pluggable vector index engine implementations. - -mod usearch; - -use store_api::storage::{VectorIndexEngine, VectorIndexEngineType}; -pub use usearch::UsearchEngine; - -use crate::vector::VectorDistanceMetric; -use crate::vector::error::Result; - -/// Configuration for creating a vector index. -#[derive(Debug, Clone)] -pub struct VectorIndexConfig { - /// The vector index engine type. - pub engine: VectorIndexEngineType, - /// The dimension of vectors in this column. - pub dim: usize, - /// The distance metric (e.g., L2sq, Cosine, InnerProduct). - pub distance_metric: VectorDistanceMetric, - /// HNSW connectivity parameter (M in the paper). - /// Higher values give better recall but use more memory. - pub connectivity: usize, - /// Expansion factor during index construction (ef_construction). - pub expansion_add: usize, - /// Expansion factor during search (ef_search). - pub expansion_search: usize, -} - -/// Creates a new vector index engine based on the engine type. -pub fn create_engine( - engine_type: VectorIndexEngineType, - config: &VectorIndexConfig, -) -> Result> { - match engine_type { - VectorIndexEngineType::Usearch => Ok(Box::new(UsearchEngine::create(config)?)), - } -} - -/// Loads a vector index engine from serialized data. -#[allow(unused)] -pub fn load_engine( - engine_type: VectorIndexEngineType, - config: &VectorIndexConfig, - data: &[u8], -) -> Result> { - match engine_type { - VectorIndexEngineType::Usearch => Ok(Box::new(UsearchEngine::load(config, data)?)), - } -} diff --git a/src/index/src/vector/engine/usearch.rs b/src/index/src/vector/engine/usearch.rs deleted file mode 100644 index 6c4352bb032..00000000000 --- a/src/index/src/vector/engine/usearch.rs +++ /dev/null @@ -1,213 +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. - -//! USearch HNSW implementation of VectorIndexEngine. - -use common_error::ext::BoxedError; -use store_api::storage::{VectorIndexEngine, VectorSearchMatches}; -use usearch::{Index, IndexOptions, ScalarKind}; - -use crate::vector::distance_metric_to_usearch; -use crate::vector::engine::VectorIndexConfig; -use crate::vector::error::{EngineSnafu, Result}; - -type EngineResult = std::result::Result; - -/// USearch-based vector index engine using HNSW algorithm. -pub struct UsearchEngine { - index: Index, -} - -impl UsearchEngine { - /// Creates a new USearch engine with the given configuration. - pub fn create(config: &VectorIndexConfig) -> Result { - let options = IndexOptions { - dimensions: config.dim, - metric: distance_metric_to_usearch(config.distance_metric), - quantization: ScalarKind::F32, - connectivity: config.connectivity, - expansion_add: config.expansion_add, - expansion_search: config.expansion_search, - multi: false, - }; - - let index = Index::new(&options).map_err(|e| { - EngineSnafu { - reason: format!("Failed to create USearch index: {}", e), - } - .build() - })?; - - Ok(Self { index }) - } - - /// Loads a USearch engine from serialized data. - #[allow(unused)] - pub fn load(config: &VectorIndexConfig, data: &[u8]) -> Result { - let options = IndexOptions { - dimensions: config.dim, - metric: distance_metric_to_usearch(config.distance_metric), - quantization: ScalarKind::F32, - // These will be loaded from serialized data - connectivity: 0, - expansion_add: 0, - expansion_search: 0, - multi: false, - }; - - let index = Index::new(&options).map_err(|e| { - EngineSnafu { - reason: format!("Failed to create USearch index for loading: {}", e), - } - .build() - })?; - - index.load_from_buffer(data).map_err(|e| { - EngineSnafu { - reason: format!("Failed to load USearch index from buffer: {}", e), - } - .build() - })?; - - Ok(Self { index }) - } - - /// Helper to create a BoxedError from a reason string. - fn boxed_engine_error(reason: String) -> BoxedError { - BoxedError::new(EngineSnafu { reason }.build()) - } -} - -impl VectorIndexEngine for UsearchEngine { - fn add(&mut self, key: u64, vector: &[f32]) -> EngineResult<()> { - // Reserve capacity if needed - if self.index.size() >= self.index.capacity() { - let new_capacity = std::cmp::max(1, self.index.capacity() * 2); - self.index.reserve(new_capacity).map_err(|e| { - Self::boxed_engine_error(format!("Failed to reserve capacity: {}", e)) - })?; - } - - self.index - .add(key, vector) - .map_err(|e| Self::boxed_engine_error(format!("Failed to add vector: {}", e))) - } - - fn search(&self, query: &[f32], k: usize) -> EngineResult { - let matches = self - .index - .search(query, k) - .map_err(|e| Self::boxed_engine_error(format!("Failed to search: {}", e)))?; - - Ok(VectorSearchMatches { - keys: matches.keys, - distances: matches.distances, - }) - } - - fn serialized_length(&self) -> usize { - self.index.serialized_length() - } - - fn save_to_buffer(&self, buffer: &mut [u8]) -> EngineResult<()> { - self.index - .save_to_buffer(buffer) - .map_err(|e| Self::boxed_engine_error(format!("Failed to save to buffer: {}", e))) - } - - fn reserve(&mut self, capacity: usize) -> EngineResult<()> { - self.index - .reserve(capacity) - .map_err(|e| Self::boxed_engine_error(format!("Failed to reserve: {}", e))) - } - - fn size(&self) -> usize { - self.index.size() - } - - fn capacity(&self) -> usize { - self.index.capacity() - } - - fn memory_usage(&self) -> usize { - self.index.memory_usage() - } -} - -#[cfg(test)] -mod tests { - use store_api::storage::VectorIndexEngineType; - - use super::*; - use crate::vector::VectorDistanceMetric; - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 4, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[test] - fn test_usearch_engine_create() { - let config = test_config(); - let engine = UsearchEngine::create(&config).unwrap(); - assert_eq!(engine.size(), 0); - } - - #[test] - fn test_usearch_engine_add_and_search() { - let config = test_config(); - let mut engine = UsearchEngine::create(&config).unwrap(); - - // Add some vectors - engine.add(0, &[1.0, 0.0, 0.0, 0.0]).unwrap(); - engine.add(1, &[0.0, 1.0, 0.0, 0.0]).unwrap(); - engine.add(2, &[0.0, 0.0, 1.0, 0.0]).unwrap(); - - assert_eq!(engine.size(), 3); - - // Search - let matches = engine.search(&[1.0, 0.0, 0.0, 0.0], 2).unwrap(); - assert_eq!(matches.keys.len(), 2); - // First result should be the exact match (key 0) - assert_eq!(matches.keys[0], 0); - } - - #[test] - fn test_usearch_engine_serialization() { - let config = test_config(); - let mut engine = UsearchEngine::create(&config).unwrap(); - - engine.add(0, &[1.0, 0.0, 0.0, 0.0]).unwrap(); - engine.add(1, &[0.0, 1.0, 0.0, 0.0]).unwrap(); - - // Serialize - let len = engine.serialized_length(); - let mut buffer = vec![0u8; len]; - engine.save_to_buffer(&mut buffer).unwrap(); - - // Load - let loaded = UsearchEngine::load(&config, &buffer).unwrap(); - assert_eq!(loaded.size(), 2); - - // Verify search works on loaded index - let matches = loaded.search(&[1.0, 0.0, 0.0, 0.0], 1).unwrap(); - assert_eq!(matches.keys[0], 0); - } -} diff --git a/src/index/src/vector/error.rs b/src/index/src/vector/error.rs deleted file mode 100644 index 10839db2068..00000000000 --- a/src/index/src/vector/error.rs +++ /dev/null @@ -1,152 +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 common_error::ext::{BoxedError, ErrorExt}; -use common_error::status_code::StatusCode; -use common_macro::stack_trace_debug; -use snafu::{Location, Snafu}; - -#[derive(Snafu)] -#[snafu(visibility(pub))] -#[stack_trace_debug] -pub enum Error { - #[snafu(display("Failed to decode protobuf"))] - DecodeProto { - #[snafu(source)] - error: prost::DecodeError, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to deserialize bitmap"))] - DeserializeBitmap { - #[snafu(source)] - error: std::io::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to serialize bitmap"))] - SerializeBitmap { - #[snafu(source)] - error: std::io::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Vector index blob too small"))] - BlobTooSmall { - min_size: usize, - actual_size: usize, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Vector index blob truncated"))] - BlobTruncated { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Invalid vector index blob"))] - InvalidBlob { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Unknown vector index engine type: {engine_type}"))] - UnknownEngineType { - engine_type: i32, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Unknown vector distance metric: {metric}"))] - UnknownDistanceMetric { - metric: i32, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Vector index engine error: {reason}"))] - Engine { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to map HNSW key to row offset: {reason}"))] - KeyMapping { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("External error"))] - External { - #[snafu(source)] - error: BoxedError, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to finish vector index: {reason}"))] - Finish { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Row count overflow: {current} + {increment} exceeds u32::MAX"))] - RowCountOverflow { - current: u64, - increment: u64, - #[snafu(implicit)] - location: Location, - }, -} - -impl ErrorExt for Error { - fn status_code(&self) -> StatusCode { - use Error::*; - - match self { - DecodeProto { .. } - | DeserializeBitmap { .. } - | SerializeBitmap { .. } - | BlobTooSmall { .. } - | BlobTruncated { .. } - | InvalidBlob { .. } - | Engine { .. } - | KeyMapping { .. } - | Finish { .. } => StatusCode::Unexpected, - - UnknownEngineType { .. } | UnknownDistanceMetric { .. } | RowCountOverflow { .. } => { - StatusCode::InvalidArguments - } - - External { error, .. } => error.status_code(), - } - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -pub type Result = std::result::Result; diff --git a/src/index/src/vector/format.rs b/src/index/src/vector/format.rs deleted file mode 100644 index 28a7f9744e9..00000000000 --- a/src/index/src/vector/format.rs +++ /dev/null @@ -1,200 +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. - -//! Vector index blob format helpers. -//! -//! Blob layout (footer-last): -//! ```text -//! +----------------------+ -//! | Null Bitmap | <- RoaringBitmap serialized, size = meta.null_bitmap_size -//! +----------------------+ -//! | Index Data | <- engine-specific serialized, size = meta.index_size -//! +----------------------+ -//! | Meta (protobuf) | <- VectorIndexMeta encoded -//! +----------------------+ -//! | Meta Size (4B) | <- u32 LE -//! +----------------------+ -//! ``` - -use greptime_proto::v1::index::{ - VectorDistanceMetric as ProtoVectorDistanceMetric, VectorIndexEngine as ProtoVectorIndexEngine, -}; -use store_api::storage::{VectorDistanceMetric, VectorIndexEngineType}; - -/// Size of the meta length field in bytes (u32 LE). -pub const META_SIZE_LEN: usize = 4; - -/// Converts Rust `VectorIndexEngineType` to proto `VectorIndexEngine`. -pub fn engine_type_to_proto(engine: VectorIndexEngineType) -> ProtoVectorIndexEngine { - match engine { - VectorIndexEngineType::Usearch => ProtoVectorIndexEngine::Usearch, - } -} - -/// Converts proto `VectorIndexEngine` to Rust `VectorIndexEngineType`. -pub fn engine_type_from_proto(engine: ProtoVectorIndexEngine) -> Option { - match engine { - ProtoVectorIndexEngine::Usearch => Some(VectorIndexEngineType::Usearch), - } -} - -/// Converts Rust `VectorDistanceMetric` to proto `VectorDistanceMetric`. -pub fn distance_metric_to_proto(metric: VectorDistanceMetric) -> ProtoVectorDistanceMetric { - match metric { - VectorDistanceMetric::L2sq => ProtoVectorDistanceMetric::L2sq, - VectorDistanceMetric::Cosine => ProtoVectorDistanceMetric::Cosine, - VectorDistanceMetric::InnerProduct => ProtoVectorDistanceMetric::InnerProduct, - } -} - -/// Converts proto `VectorDistanceMetric` to Rust `VectorDistanceMetric`. -pub fn distance_metric_from_proto( - metric: ProtoVectorDistanceMetric, -) -> Option { - match metric { - ProtoVectorDistanceMetric::L2sq => Some(VectorDistanceMetric::L2sq), - ProtoVectorDistanceMetric::Cosine => Some(VectorDistanceMetric::Cosine), - ProtoVectorDistanceMetric::InnerProduct => Some(VectorDistanceMetric::InnerProduct), - } -} - -#[cfg(test)] -mod tests { - use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; - use prost::Message; - use roaring::RoaringBitmap; - - use super::*; - - #[test] - fn test_engine_type_roundtrip() { - let engine = VectorIndexEngineType::Usearch; - let proto = engine_type_to_proto(engine); - let back = engine_type_from_proto(proto).unwrap(); - assert_eq!(engine, back); - } - - #[test] - fn test_distance_metric_roundtrip() { - for metric in [ - VectorDistanceMetric::L2sq, - VectorDistanceMetric::Cosine, - VectorDistanceMetric::InnerProduct, - ] { - let proto = distance_metric_to_proto(metric); - let back = distance_metric_from_proto(proto).unwrap(); - assert_eq!(metric, back); - } - } - - #[test] - fn test_vector_index_meta_encode_decode() { - let meta = VectorIndexMeta { - engine: engine_type_to_proto(VectorIndexEngineType::Usearch).into(), - dim: 128, - metric: distance_metric_to_proto(VectorDistanceMetric::Cosine).into(), - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - null_bitmap_size: 10, - index_size: 1024, - stats: Some(VectorIndexStats { - total_row_count: 100, - indexed_row_count: 95, - null_count: 5, - }), - }; - - let encoded = meta.encode_to_vec(); - let decoded = VectorIndexMeta::decode(encoded.as_slice()).unwrap(); - - assert_eq!(decoded.engine, meta.engine); - assert_eq!(decoded.dim, meta.dim); - assert_eq!(decoded.metric, meta.metric); - assert_eq!(decoded.connectivity, meta.connectivity); - assert_eq!(decoded.expansion_add, meta.expansion_add); - assert_eq!(decoded.expansion_search, meta.expansion_search); - assert_eq!(decoded.null_bitmap_size, meta.null_bitmap_size); - assert_eq!(decoded.index_size, meta.index_size); - - let stats = decoded.stats.unwrap(); - assert_eq!(stats.total_row_count, 100); - assert_eq!(stats.indexed_row_count, 95); - assert_eq!(stats.null_count, 5); - } - - #[test] - fn test_blob_format_roundtrip() { - // Create test data - let null_bitmap = RoaringBitmap::from_iter([1u32, 3, 5]); - let mut null_bitmap_bytes = Vec::new(); - null_bitmap.serialize_into(&mut null_bitmap_bytes).unwrap(); - - let index_data = b"fake_index_data_here"; - - let meta = VectorIndexMeta { - engine: engine_type_to_proto(VectorIndexEngineType::Usearch).into(), - dim: 4, - metric: distance_metric_to_proto(VectorDistanceMetric::L2sq).into(), - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_data.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: 10, - indexed_row_count: 7, - null_count: 3, - }), - }; - - // Build blob: null_bitmap | index_data | meta | meta_size - let mut blob = Vec::new(); - blob.extend_from_slice(&null_bitmap_bytes); - blob.extend_from_slice(index_data); - let meta_bytes = meta.encode_to_vec(); - blob.extend_from_slice(&meta_bytes); - let meta_size = meta_bytes.len() as u32; - blob.extend_from_slice(&meta_size.to_le_bytes()); - - // Parse blob - let blob_len = blob.len(); - let meta_size_read = - u32::from_le_bytes(blob[blob_len - META_SIZE_LEN..].try_into().unwrap()) as usize; - assert_eq!(meta_size_read, meta_bytes.len()); - - let meta_start = blob_len - META_SIZE_LEN - meta_size_read; - let decoded_meta = - VectorIndexMeta::decode(&blob[meta_start..blob_len - META_SIZE_LEN]).unwrap(); - - assert_eq!( - decoded_meta.null_bitmap_size, - null_bitmap_bytes.len() as u32 - ); - assert_eq!(decoded_meta.index_size, index_data.len() as u64); - - // Extract null bitmap - let null_bitmap_end = decoded_meta.null_bitmap_size as usize; - let restored_bitmap = RoaringBitmap::deserialize_from(&blob[..null_bitmap_end]).unwrap(); - assert_eq!(restored_bitmap.len(), 3); - assert!(restored_bitmap.contains(1)); - assert!(restored_bitmap.contains(3)); - assert!(restored_bitmap.contains(5)); - - // Extract index data - let index_start = null_bitmap_end; - let index_end = index_start + decoded_meta.index_size as usize; - assert_eq!(&blob[index_start..index_end], index_data); - } -} diff --git a/src/index/src/vector/util.rs b/src/index/src/vector/util.rs deleted file mode 100644 index 7610a07e7cc..00000000000 --- a/src/index/src/vector/util.rs +++ /dev/null @@ -1,108 +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. - -//! Utility functions for vector index operations. - -use std::borrow::Cow; - -/// Converts a byte slice (little-endian format) to f32 slice, handling unaligned data gracefully. -/// Returns `Cow::Borrowed` for aligned data on little-endian systems (zero-copy) -/// or `Cow::Owned` for unaligned data or big-endian systems. -/// -/// # Panics -/// -/// Panics if the byte slice length is not a multiple of 4. -pub fn bytes_to_f32_slice(bytes: &[u8]) -> Cow<'_, [f32]> { - assert!( - bytes.len().is_multiple_of(4), - "Vector bytes length {} is not a multiple of 4", - bytes.len() - ); - - if bytes.is_empty() { - return Cow::Borrowed(&[]); - } - - let ptr = bytes.as_ptr(); - // Fast path: zero-copy only when data is aligned AND we're on little-endian system - // (since vector data is stored in little-endian format) - #[cfg(target_endian = "little")] - if (ptr as usize).is_multiple_of(std::mem::align_of::()) { - // Safety: We've verified alignment and length requirements, - // and on little-endian systems the byte representation matches f32 layout - return Cow::Borrowed(unsafe { - std::slice::from_raw_parts(ptr as *const f32, bytes.len() / 4) - }); - } - - // Slow path: data is not aligned or we're on big-endian system - let floats: Vec = bytes - .chunks_exact(4) - .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect(); - Cow::Owned(floats) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_bytes_to_f32_slice() { - let floats = [1.0f32, 2.0, 3.0, 4.0]; - let bytes: Vec = floats.iter().flat_map(|f| f.to_le_bytes()).collect(); - - let result = bytes_to_f32_slice(&bytes); - assert_eq!(result.len(), 4); - assert_eq!(result[0], 1.0); - assert_eq!(result[1], 2.0); - assert_eq!(result[2], 3.0); - assert_eq!(result[3], 4.0); - } - - #[test] - fn test_bytes_to_f32_slice_unaligned() { - // Create a buffer with an extra byte at the start to force misalignment - let floats = [1.0f32, 2.0, 3.0, 4.0]; - let mut bytes: Vec = vec![0u8]; // padding byte - bytes.extend(floats.iter().flat_map(|f| f.to_le_bytes())); - - // Take a slice starting at offset 1 (unaligned) - let unaligned_bytes = &bytes[1..]; - - // Verify it's actually unaligned - let ptr = unaligned_bytes.as_ptr(); - let is_aligned = (ptr as usize).is_multiple_of(std::mem::align_of::()); - - // The function should work regardless of alignment - let result = bytes_to_f32_slice(unaligned_bytes); - assert_eq!(result.len(), 4); - assert_eq!(result[0], 1.0); - assert_eq!(result[1], 2.0); - assert_eq!(result[2], 3.0); - assert_eq!(result[3], 4.0); - - // If it was unaligned, it should return an owned Vec (Cow::Owned) - if !is_aligned { - assert!(matches!(result, Cow::Owned(_))); - } - } - - #[test] - fn test_bytes_to_f32_slice_empty() { - let bytes: &[u8] = &[]; - let result = bytes_to_f32_slice(bytes); - assert!(result.is_empty()); - } -} diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 15560b06b2c..7e4da1557e1 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -11,6 +11,9 @@ protoc-rust.workspace = true [lints] workspace = true +[features] +testing = [] + [dependencies] async-stream.workspace = true async-trait.workspace = true @@ -31,7 +34,6 @@ delta-encoding = "0.4" derive_builder.workspace = true futures.workspace = true futures-util.workspace = true -itertools.workspace = true lazy_static.workspace = true object-store.workspace = true pin-project.workspace = true @@ -49,8 +51,13 @@ tokio.workspace = true common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true common-wal = { workspace = true, features = ["testing"] } -itertools.workspace = true object-store = { workspace = true, features = ["testing"] } rand.workspace = true rskafka = { workspace = true, features = ["unstable-fuzzing"] } uuid.workspace = true + +[package.metadata.cargo-shear] +# `protobuf` is required by the code generated by `build.rs`: the generated +# protobuf bindings (included via `include!`) reference `::protobuf` paths, +# but no source file does. +ignored = ["protobuf"] diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index de304e192d2..85000dbea3c 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -350,6 +350,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Incomplete multipart WAL entry of region {}", region_id))] + IncompleteWalEntry { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("WAL object sequence is exhausted, last sequence: {}", last_object_seq))] WalObjectSequenceExhausted { last_object_seq: u64, @@ -357,6 +364,16 @@ pub enum Error { location: Location, }, + #[snafu(display( + "WAL object sequence {} is not settled: the open batch has assigned entry ids under it", + object_seq + ))] + WalObjectSequenceUnsettled { + object_seq: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "WAL entry positions of region {} in one object are exhausted", region_id @@ -374,6 +391,32 @@ pub enum Error { location: Location, }, + #[snafu(display( + "WAL object {} was written by the earlier epoch {}, this store writes epoch {}", + path, + existing_epoch, + epoch + ))] + StaleWalObject { + path: String, + existing_epoch: u64, + epoch: u64, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "WAL object {} that starts epoch {} was already present, so the open cannot tell whether it wrote it", + path, + epoch + ))] + UnconfirmedWalEpochStart { + path: String, + epoch: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to {} WAL object, path: {}", operation, path))] WalObjectStore { operation: &'static str, @@ -405,16 +448,15 @@ pub enum Error { location: Location, }, - #[snafu(display("Object store WAL operation failed"))] - ObjectStoreWal { - source: Arc, + #[snafu(display("Object store WAL log store is stopped"))] + ObjectStoreWalStopped { #[snafu(implicit)] location: Location, }, - /// Appending to the object store WAL is unsupported. - #[snafu(display("Object store WAL operation is not supported"))] - UnsupportedObjectStoreWalOperation { + #[snafu(display("Object store WAL operation failed"))] + ObjectStoreWal { + source: Arc, #[snafu(implicit)] location: Location, }, @@ -452,6 +494,7 @@ impl ErrorExt for Error { | InvalidWalObjectStore { .. } | MismatchedWalPrefix { .. } | MismatchedWalRegion { .. } + | IncompleteWalEntry { .. } | InvalidWalEntryRange { .. } => StatusCode::InvalidArguments, StartWalTask { .. } | StopWalTask { .. } @@ -465,14 +508,15 @@ impl ErrorExt for Error { | OrderedBatchProducerStopped { .. } | WaitProduceResultReceiver { .. } | WaitDumpIndex { .. } - | MetaLengthExceededLimit { .. } => StatusCode::Internal, + | MetaLengthExceededLimit { .. } + | ObjectStoreWalStopped { .. } => StatusCode::Internal, CorruptedWalObject { .. } | WalObjectConflict { .. } | WalObjectSequenceExhausted { .. } | WalEntryPositionExhausted { .. } => StatusCode::Unexpected, + WalObjectSequenceUnsettled { .. } => StatusCode::IllegalState, - UnsupportedObjectStoreWalOperation { .. } => StatusCode::Unsupported, InvalidWalObject { source, .. } => source.status_code(), ObjectStoreWal { source, .. } => source.status_code(), @@ -481,6 +525,8 @@ impl ErrorExt for Error { | WriteIndex { .. } | ReadIndex { .. } | WalObjectStore { .. } + | StaleWalObject { .. } + | UnconfirmedWalEpochStart { .. } | Io { .. } => StatusCode::StorageUnavailable, // Raft engine FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } => { @@ -513,7 +559,12 @@ impl ErrorExt for Error { | WalObjectStore { error, .. } => retry_hint_from_opendal_error(error), ObjectStoreWal { source, .. } => source.retry_hint(), Io { error, .. } => retry_hint_from_io_error(error), - FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } => RetryHint::Retryable, + FetchEntry { .. } + | RaftEngine { .. } + | AddEntryLogBatch { .. } + | WalObjectSequenceUnsettled { .. } + | StaleWalObject { .. } + | UnconfirmedWalEpochStart { .. } => RetryHint::Retryable, ProduceRecord { error, .. } => match error { rskafka::client::producer::Error::Client(error) => { rskafka_client_error_to_retry_hint(error) diff --git a/src/log-store/src/kafka/consumer.rs b/src/log-store/src/kafka/consumer.rs index 8a7cba5a4b6..d064a9bc9ac 100644 --- a/src/log-store/src/kafka/consumer.rs +++ b/src/log-store/src/kafka/consumer.rs @@ -158,13 +158,10 @@ impl RecordsBuffer { impl RecordsBuffer { fn pop_front(&mut self) -> Option { while let Some(index) = self.index.peek() { - if let Some(record_and_offset) = self.buffer.pop_front() { - if index == record_and_offset.offset as u64 { - self.index.next(); - return Some(record_and_offset); - } - } else { - return None; + let record_and_offset = self.buffer.pop_front()?; + if index == record_and_offset.offset as u64 { + self.index.next(); + return Some(record_and_offset); } } diff --git a/src/log-store/src/lib.rs b/src/log-store/src/lib.rs index 5773efa7aaa..0cc133d81e7 100644 --- a/src/log-store/src/lib.rs +++ b/src/log-store/src/lib.rs @@ -19,3 +19,5 @@ pub mod noop; pub mod object_store_wal; pub mod raft_engine; pub mod test_util; + +pub use object_store_wal::ObjectStoreLogStore; diff --git a/src/log-store/src/metrics.rs b/src/log-store/src/metrics.rs index f12ad3e815a..03a3dc1f174 100644 --- a/src/log-store/src/metrics.rs +++ b/src/log-store/src/metrics.rs @@ -89,4 +89,193 @@ lazy_static! { &[LOGSTORE_LABEL, PARTITION_LABEL], ) .unwrap(); + + /// Counter of segments a read of the object store logstore skipped because + /// they did not decode. + pub static ref METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_skipped_segments_total", + "object store logstore skipped corrupted segments total", + ) + .unwrap(); + /// Counter of objects the object store logstore deleted because no + /// region needs them any more. + pub static ref METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_deleted_objects_total", + "object store logstore deleted objects total", + ) + .unwrap(); + /// Counter of object deletes of the object store logstore that failed and + /// are attempted again at the next collection. + pub static ref METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_failed_deletes_total", + "object store logstore failed object deletes total", + ) + .unwrap(); + + /// How long an object of the object store logstore takes to become durable + /// once its batch is sealed, sampled for every conditional create that + /// succeeds, an identical retry included. The time runs from the seal, so + /// it includes any wait for a create slot and, in the `enqueued` mode, + /// every repeated attempt and the delay before it. + /// + /// A create is a few milliseconds on a local object store and tens of + /// milliseconds on a cloud one, so the buckets start at 1ms and reach 8s. + pub static ref METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_seal_to_durable_seconds", + "object store logstore seconds from sealing a batch to the success of its conditional create, slot waits and retries included", + exponential_buckets(0.001, 2.0, 14).unwrap(), + ) + .unwrap(); + + /// How long an append of the object store logstore waits for its + /// acknowledgement in the `durable` mode, which is the write latency a + /// caller sees. + /// + /// The wait is one seal plus every earlier object plus one create, so with + /// the default 100ms flush interval it sits in the tens of milliseconds; + /// the buckets straddle that and leave room for a stalled uploader. + pub static ref METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_append_ack_seconds", + "object store logstore seconds an append waits for its acknowledgement in the durable mode", + exponential_buckets(0.002, 2.0, 14).unwrap(), + ) + .unwrap(); + + /// How large the objects of the object store logstore are, which says + /// whether batches seal on size or on the flush interval. + /// + /// The buckets are powers of two from 1KiB, the size of a timer-sealed + /// batch of one small entry, to 16MiB, twice the 8MiB default batch size, + /// which is a threshold rather than an upper bound. + pub static ref METRIC_OBJECT_STORE_WAL_OBJECT_BYTES: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_object_bytes", + "object store logstore bytes of each object sealed", + exponential_buckets(1024.0, 2.0, 15).unwrap(), + ) + .unwrap(); + + /// How many entries the objects of the object store logstore hold, which + /// says how many writes one request amortises. + /// + /// The largest finite bucket is 2^20, the position limit of one region in + /// one object. An object holds the entries of many regions, so its total + /// can pass that; such an object falls in the overflow bucket. + pub static ref METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_object_entries", + "object store logstore entries of each object sealed", + exponential_buckets(1.0, 4.0, 11).unwrap(), + ) + .unwrap(); + + /// Counter of conditional creates of the object store logstore that + /// succeeded, an identical retry that found its own object included. A + /// create that failed is counted in the create failure or conflict counter + /// instead. + pub static ref METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_created_objects_total", + "object store logstore conditional creates that succeeded, identical retries included", + ) + .unwrap(); + + /// Number of objects the object store logstore has indexed, which is what + /// a recovery of the prefix would have to read. + pub static ref METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS: IntGauge = register_int_gauge!( + "greptime_logstore_object_store_wal_indexed_objects", + "object store logstore objects currently indexed", + ) + .unwrap(); + + /// Bytes of the objects the object store logstore has indexed, which is + /// what the WAL occupies under its prefix. + pub static ref METRIC_OBJECT_STORE_WAL_INDEXED_BYTES: IntGauge = register_int_gauge!( + "greptime_logstore_object_store_wal_indexed_bytes", + "object store logstore bytes currently indexed", + ) + .unwrap(); + + /// How long a recovery of the object store logstore that succeeded took to + /// rebuild its catalog, which is the part of a restart the WAL is + /// responsible for. A recovery that was abandoned records nothing. + /// + /// Recovery reads a few small ranges per object, so the buckets run from + /// 10ms, an empty prefix, to minutes on a prefix with many objects. + pub static ref METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_recovery_seconds", + "object store logstore seconds spent by a recovery that succeeded", + exponential_buckets(0.01, 3.0, 10).unwrap(), + ) + .unwrap(); + + /// Counter of objects a recovery of the object store logstore that + /// succeeded read footers from, which is what its time is spent on. It + /// counts the objects of the same recoveries the timer above measures, so + /// an abandoned one contributes to neither. + pub static ref METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_recovered_objects_total", + "object store logstore objects read by a recovery that succeeded total", + ) + .unwrap(); + + /// How long a read stream of the object store logstore took to serve the + /// replay of one region, recorded once the stream has been read to its + /// end; one that fails or is dropped part way records nothing. + /// + /// A read fetches one segment per object of the region, so the buckets run + /// from 5ms, a region with nothing to replay, to minutes. + pub static ref METRIC_OBJECT_STORE_WAL_READ_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_read_seconds", + "object store logstore seconds spent serving a region read that ran to its end", + exponential_buckets(0.005, 3.0, 10).unwrap(), + ) + .unwrap(); + + /// Counter of conditional creates of the object store logstore that found + /// an object with different content. In the `durable` mode an object of an + /// earlier epoch fails the batch like a transient failure; any other + /// conflict poisons the store. + pub static ref METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_create_conflicts_total", + "object store logstore conflicting conditional creates total", + ) + .unwrap(); + + /// Counter of conditional creates of the object store logstore the object + /// store did not confirm, which fail the appends behind them in the + /// `durable` mode and are repeated in the `enqueued` mode. + pub static ref METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_create_failures_total", + "object store logstore transient conditional create failures total", + ) + .unwrap(); + + /// Counter of object store logstores that hit an error they cannot recover + /// from, after which every operation of the store fails until it reopens. + pub static ref METRIC_OBJECT_STORE_WAL_POISONED_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_poisoned_total", + "object store logstores poisoned total", + ) + .unwrap(); + + /// Counter of appends the object store logstore held back because the + /// unpersisted backlog reached a threshold in the `enqueued` mode, which + /// says the write rate is above what the uploads sustain. + pub static ref METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_stalled_appends_total", + "object store logstore appends stalled by the backlog thresholds total", + ) + .unwrap(); + + /// How long such an append stays held back, until it is admitted or until + /// it is refused, a stop, a terminal error or the teardown of the store + /// releasing it included, which is the latency the `enqueued` mode pays + /// once the backlog is at a threshold. + /// + /// A stalled append waits for one upload to complete, so the buckets run + /// from 5ms to half a minute of a struggling object store. + pub static ref METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_stalled_append_seconds", + "object store logstore seconds a stalled append is held back before it is admitted or refused", + exponential_buckets(0.005, 3.0, 9).unwrap(), + ) + .unwrap(); } diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index 70bc39dedca..6a25afb3e77 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -23,39 +23,52 @@ //! ``` //! //! The header carries the magic `GTWALOBJ`, the format version, the object -//! sequence and the instance that wrote the object. Each segment holds the -//! entries of exactly one region, ordered by entry id, and segments are ordered -//! by region id. The footer indexes every segment with its region id, entry id +//! sequence, the epoch of the instance that wrote the object, the link to the +//! object it extends, the reclaim boundary, and a CRC32 of the header. Each +//! segment holds the entries of exactly one region, ordered by entry id, and +//! segments are ordered by region id. The footer indexes every segment with its region id, entry id //! range, byte range and CRC32. The fixed-size trailer points at the footer and //! carries the CRC32 of the footer and of the whole object, so a reader locates //! the footer by reading the fixed-length trailer at the end of the object. //! //! Object sequences increase monotonically within one prefix and may leave -//! gaps, so recovery continues after the largest sequence it indexed. An object -//! is created conditionally: rewriting a sequence with the content it already +//! gaps, so recovery continues after the largest sequence present. An object is +//! created conditionally: rewriting a sequence with the content it already //! holds is a no-op at the object store, while different content under a taken -//! sequence is a conflict. Recovery lists the objects, reads and verifies only -//! the header, trailer and footer of each, and indexes the footers in sequence -//! order to rebuild the object catalog, which rejects a sequence it already -//! holds. Segments are read and checksummed only when a read decodes them. +//! sequence is a conflict. A create can fail with an unknown outcome and its +//! object can still appear, so objects form a chain: every object links to the +//! object it extends by sequence and epoch, and recovery replays only the chain +//! that ends at the complete object with the largest epoch and sequence. An +//! object is complete when every link on its chain names a present object of +//! the recorded epoch, back to an object that starts the chain or to a link +//! below the reclaim boundary of the object. Objects off the chain are +//! orphans, which are never replayed but keep their sequences. Each open +//! writes an object without segments above every present object before it +//! accepts writes, and its epoch is one above the sequence of that object, so +//! no two instances share an epoch and a late object of an earlier instance +//! never ends the chain. +//! +//! The reclaim boundary of an object states that every object with a lower +//! sequence was fully flushed when the object was sealed, so the chain need +//! not reach below it. Every present object below the boundary of the last +//! indexed object is deleted, on the chain or not. +//! +//! Recovery lists the objects, reads and verifies only the header, trailer and +//! footer of each, and indexes the footers of the chain in sequence order to +//! rebuild the object catalog, which rejects a sequence it already holds. +//! Segments are read and checksummed only when a read decodes them; a segment +//! that does not decode is fetched once more and then skipped as a hole of its +//! region or fails the read, as `on_corrupted_segment` says. //! //! Entry ids are object-sequence-major, see [`entry_id`]: the high bits of an //! id name the object that holds the entry, the low bits its position among //! the entries of its region in that object. -#[allow(dead_code)] mod batch; -#[allow(dead_code)] mod catalog; -#[allow(dead_code)] mod format; -#[allow(dead_code)] mod io; - -#[allow(dead_code)] mod store; -#[allow(unused_imports)] -pub(crate) use batch::entry_id; -#[allow(unused_imports)] -pub(crate) use store::ObjectStoreLogStore; +pub use batch::entry_id; +pub use store::{ObjectStoreLogStore, WalHole}; diff --git a/src/log-store/src/object_store_wal/batch.rs b/src/log-store/src/object_store_wal/batch.rs index b07ac8e8cb0..b1873c8a219 100644 --- a/src/log-store/src/object_store_wal/batch.rs +++ b/src/log-store/src/object_store_wal/batch.rs @@ -45,7 +45,7 @@ pub(crate) const OBJECT_SEQ_LIMIT: u64 = 1 << (u64::BITS - POSITION_BITS); /// at one so that id zero, the watermark of a region without entries, is never /// assigned. A region's ids increase with the object sequence and have gaps /// wherever other regions or other positions took the sequence. -pub(crate) fn entry_id(object_seq: u64, position: u64) -> EntryId { +pub fn entry_id(object_seq: u64, position: u64) -> EntryId { debug_assert!(object_seq < OBJECT_SEQ_LIMIT && (1..POSITION_LIMIT).contains(&position)); (object_seq << POSITION_BITS) | position } @@ -145,6 +145,10 @@ impl OpenBatch { self.entries.is_empty() } + pub(crate) fn holds_region(&self, region_id: RegionId) -> bool { + self.positions.contains_key(®ion_id) + } + /// Returns the estimated size of the admitted entries. pub(crate) fn estimated_bytes(&self) -> usize { self.estimated_bytes diff --git a/src/log-store/src/object_store_wal/catalog.rs b/src/log-store/src/object_store_wal/catalog.rs index 5e9650d14a5..70f7aa5b329 100644 --- a/src/log-store/src/object_store_wal/catalog.rs +++ b/src/log-store/src/object_store_wal/catalog.rs @@ -14,10 +14,11 @@ //! In-memory index over the footers of the objects of one WAL prefix. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::ops::Bound::{Excluded, Unbounded}; use snafu::{OptionExt, ensure}; +use store_api::logstore::EntryId; use store_api::storage::RegionId; use crate::error::{ @@ -32,25 +33,22 @@ use crate::object_store_wal::format::FooterEntry; pub(crate) struct ObjectCatalog { objects: BTreeMap>, regions: BTreeMap>, + /// Largest entry id ever indexed per region, kept after the object that + /// held it was removed. + max_entry_ids: HashMap, } impl ObjectCatalog { /// Indexes the footer of the object `object_seq`. Objects may be inserted /// in any order, which lets recovery index them as it discovers them. /// Inserting a sequence that is already indexed is rejected, whether or not - /// the footer matches the indexed one. + /// the footer matches the indexed one. An object without segments holds no + /// entries but still takes its sequence. pub(crate) fn insert_object( &mut self, object_seq: u64, mut footer: Vec, ) -> Result<()> { - ensure!( - !footer.is_empty(), - CorruptedWalObjectSnafu { - reason: format!("object {object_seq} has an empty footer"), - } - ); - footer.sort_unstable_by_key(|entry| entry.region_id); for entries in footer.windows(2) { ensure!( @@ -129,11 +127,71 @@ impl ObjectCatalog { .entry(entry.region_id) .or_default() .insert(object_seq, entry.clone()); + self.max_entry_ids + .entry(entry.region_id) + .and_modify(|max| *max = (*max).max(entry.max_entry_id)) + .or_insert(entry.max_entry_id); } self.objects.insert(object_seq, footer); Ok(()) } + /// Removes the object `object_seq` from the index once it was deleted and + /// returns the footer it held, or `None` if it was not indexed. The + /// largest entry id of every region it held is kept. + pub(crate) fn remove_object(&mut self, object_seq: u64) -> Option> { + let footer = self.objects.remove(&object_seq)?; + for entry in &footer { + if let Some(objects) = self.regions.get_mut(&entry.region_id) { + objects.remove(&object_seq); + if objects.is_empty() { + self.regions.remove(&entry.region_id); + } + } + } + Some(footer) + } + + /// Returns true while the object `object_seq` is indexed. + pub(crate) fn contains_object(&self, object_seq: u64) -> bool { + self.objects.contains_key(&object_seq) + } + + /// Returns the reclaim boundary the indexed objects and the obsolete + /// watermarks allow, given that every object below `from` holds no live + /// segment: the first indexed object at or above `from` that holds a + /// live segment, or the sequence after the last indexed object when none + /// does, and never less than `from`. + /// + /// A segment is live when its maximum entry id is above the watermark of + /// its region, or its region has no watermark; a region whose watermark + /// is [`EntryId::MAX`] holds no live segment. Watermarks never move down + /// and a region first appears in an object above the boundary, so the + /// scan starts at the previous boundary and never reads an object twice + /// once the boundary has passed it. + pub(crate) fn reclaim_boundary( + &self, + from: u64, + obsolete_entry_ids: &HashMap, + ) -> u64 { + let live = self.objects.range(from..).find(|(_, footer)| { + footer.iter().any(|entry| { + obsolete_entry_ids + .get(&entry.region_id) + .is_none_or(|obsolete| entry.max_entry_id > *obsolete) + }) + }); + match live { + Some((&object_seq, _)) => object_seq, + None => self + .objects + .last_key_value() + .map_or(from, |(&object_seq, _)| { + object_seq.saturating_add(1).max(from) + }), + } + } + /// Returns the objects that hold entries of `region_id` overlapping /// `start_entry_id..=end_entry_id`, ordered by object sequence. pub(crate) fn objects_for_entry_range( @@ -163,12 +221,10 @@ impl ObjectCatalog { .collect()) } - /// Returns the largest entry id indexed for `region_id`. + /// Returns the largest entry id ever indexed for `region_id`, which is + /// kept after the object holding it was removed. pub(crate) fn region_max_entry_id(&self, region_id: RegionId) -> Option { - self.regions - .get(®ion_id)? - .last_key_value() - .map(|(_, entry)| entry.max_entry_id) + self.max_entry_ids.get(®ion_id).copied() } /// Returns the sequence to assign to the next object written after recovery. @@ -189,9 +245,9 @@ impl ObjectCatalog { .context(WalObjectSequenceExhaustedSnafu { last_object_seq })?, }; let floor = self - .regions - .keys() - .filter_map(|region_id| self.region_max_entry_id(*region_id)) + .max_entry_ids + .values() + .copied() .map(sequence_floor) .max() .unwrap_or(0); @@ -310,6 +366,11 @@ mod tests { .unwrap(); assert_eq!(5, catalog.next_object_seq().unwrap()); + + // An object without segments takes its sequence all the same. + catalog.insert_object(7, Vec::new()).unwrap(); + assert_eq!(8, catalog.next_object_seq().unwrap()); + assert_eq!(Some(12), catalog.region_max_entry_id(region_id)); } #[test] @@ -416,11 +477,6 @@ mod tests { ), "duplicate footer entries", ); - assert_corrupted( - catalog.insert_object(1, vec![]), - "object 1 has an empty footer", - ); - let mut invalid = footer_entry(region_id, 2, 1); invalid.entry_count = 0; assert_corrupted( @@ -437,6 +493,98 @@ mod tests { ); } + #[test] + fn test_catalog_reclaim_boundary_follows_the_watermarks() { + let region_one = RegionId::new(1, 1); + let region_two = RegionId::new(2, 1); + let mut catalog = ObjectCatalog::default(); + catalog.insert_object(0, Vec::new()).unwrap(); + catalog + .insert_object( + 1, + vec![ + footer_entry(region_one, entry_id(1, 1), entry_id(1, 2)), + footer_entry(region_two, entry_id(1, 1), entry_id(1, 1)), + ], + ) + .unwrap(); + catalog + .insert_object( + 2, + vec![footer_entry(region_one, entry_id(2, 1), entry_id(2, 1))], + ) + .unwrap(); + catalog + .insert_object( + 3, + vec![footer_entry(region_two, entry_id(3, 1), entry_id(3, 1))], + ) + .unwrap(); + let boundary = |catalog: &ObjectCatalog, obsolete: &[(RegionId, EntryId)]| { + catalog.reclaim_boundary(0, &obsolete.iter().copied().collect()) + }; + + // A region without a watermark holds the boundary at its first object. + assert_eq!(1, boundary(&catalog, &[])); + assert_eq!(1, boundary(&catalog, &[(region_one, entry_id(2, 1))])); + // A watermark inside an object keeps that object. + assert_eq!( + 1, + boundary( + &catalog, + &[(region_one, entry_id(1, 1)), (region_two, EntryId::MAX)] + ) + ); + assert_eq!( + 2, + boundary( + &catalog, + &[(region_one, entry_id(1, 2)), (region_two, entry_id(1, 1))] + ) + ); + assert_eq!( + 3, + boundary( + &catalog, + &[(region_one, entry_id(2, 1)), (region_two, entry_id(1, 1))] + ) + ); + // A region of `obsolete_all` contributes nothing, and when no region + // does, the boundary is past the last object. + assert_eq!( + 4, + boundary( + &catalog, + &[(region_one, entry_id(2, 1)), (region_two, EntryId::MAX)] + ) + ); + + // Objects below the previous boundary are not scanned again. + assert_eq!(3, catalog.reclaim_boundary(3, &HashMap::new())); + assert_eq!(5, catalog.reclaim_boundary(5, &HashMap::new())); + + // Removing objects keeps the largest entry id of every region, and a + // region without indexed objects contributes nothing. + for object_seq in [0, 1, 2] { + assert!(catalog.remove_object(object_seq).is_some()); + } + assert!(catalog.remove_object(2).is_none()); + assert!(!catalog.contains_object(2)); + assert!(catalog.contains_object(3)); + assert_eq!( + Some(entry_id(2, 1)), + catalog.region_max_entry_id(region_one) + ); + assert!( + catalog + .objects_for_entry_range(region_one, 0, EntryId::MAX) + .unwrap() + .is_empty() + ); + assert_eq!(3, boundary(&catalog, &[])); + assert_eq!(4, catalog.next_object_seq().unwrap()); + } + fn footer_entry(region_id: RegionId, min_entry_id: u64, max_entry_id: u64) -> FooterEntry { FooterEntry { region_id, diff --git a/src/log-store/src/object_store_wal/format.rs b/src/log-store/src/object_store_wal/format.rs index 62fe290aa82..7a255711784 100644 --- a/src/log-store/src/object_store_wal/format.rs +++ b/src/log-store/src/object_store_wal/format.rs @@ -27,9 +27,14 @@ use crate::error::{CorruptedWalObjectSnafu, Result}; const HEADER_MAGIC: &[u8; 8] = b"GTWALOBJ"; const TRAILER_MAGIC: &[u8; 8] = b"GTWALTRL"; const FORMAT_VERSION: u16 = 1; +/// Predecessor sequence recorded for an object that starts a chain. Object +/// sequences stay below `OBJECT_SEQ_LIMIT`, so it never names an object. +const NO_PREDECESSOR: u64 = u64::MAX; -/// Length of the object header: magic, version, object sequence, writer instance. -pub(crate) const HEADER_LEN: usize = 8 + 2 + 8 + 16; +/// Length of the object header: magic, version, object sequence, epoch, +/// predecessor sequence, predecessor epoch, reclaim boundary and the CRC32 of +/// the header. +pub(crate) const HEADER_LEN: usize = 8 + 2 + 8 + 8 + 8 + 8 + 8 + 4; /// Length of the fixed trailer: footer offset, footer length, footer CRC32, /// object CRC32 and magic. pub(crate) const TRAILER_LEN: usize = 8 + 8 + 4 + 4 + 8; @@ -44,7 +49,24 @@ const FOOTER_COUNT_LEN: usize = 4; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct Header { pub(crate) object_seq: u64, - pub(crate) writer_instance: [u8; 16], + /// Epoch of the store instance that wrote the object: one above the + /// sequence of the start object the instance created, so the epoch also + /// identifies the instance. + pub(crate) epoch: u64, + /// The object this object extends, `None` for the first object of a chain. + pub(crate) prev: Option, + /// Every object below this sequence was fully flushed when this object + /// was sealed, so the chain recovery replays need not reach below it. It + /// is at or below the sequence of the object. + pub(crate) reclaim_boundary: u64, +} + +/// Names the object another object extends: its sequence and the epoch +/// recorded in its header. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ChainLink { + pub(crate) object_seq: u64, + pub(crate) epoch: u64, } /// A single WAL entry inside an object. @@ -94,15 +116,9 @@ pub(crate) struct DecodedObject { } /// Encodes `records` into one object. Records are grouped into a segment per -/// region; entry ids within a region must be unique. +/// region; entry ids within a region must be unique. An object without records +/// only extends the chain. pub(crate) fn encode_object(header: Header, records: &[Record]) -> Result { - ensure!( - !records.is_empty(), - CorruptedWalObjectSnafu { - reason: "object has no records", - } - ); - let mut grouped = BTreeMap::>::new(); for record in records { grouped.entry(record.region_id).or_default().push(record); @@ -180,7 +196,8 @@ pub(crate) fn object_len(footer: &[FooterEntry]) -> u64 { segments + framing as u64 } -/// Decodes the header from the first [`HEADER_LEN`] bytes of an object. +/// Decodes and verifies the header from the first [`HEADER_LEN`] bytes of an +/// object. pub(crate) fn decode_header(bytes: &[u8]) -> Result
{ ensure!( bytes.len() >= HEADER_LEN, @@ -207,13 +224,35 @@ pub(crate) fn decode_header(bytes: &[u8]) -> Result
{ } ); let object_seq = reader.u64("header")?; - let writer_instance = reader - .take(16, "header")? - .try_into() - .expect("writer instance has a fixed length"); + let epoch = reader.u64("header")?; + let prev_object_seq = reader.u64("header")?; + let prev_epoch = reader.u64("header")?; + let reclaim_boundary = reader.u64("header")?; + let header_crc32 = reader.u32("header")?; + let checksum = crc32fast::hash(&bytes[..HEADER_LEN - 4]); + ensure!( + checksum == header_crc32, + CorruptedWalObjectSnafu { + reason: checksum_mismatch("header", header_crc32, checksum), + } + ); + ensure!( + reclaim_boundary <= object_seq, + CorruptedWalObjectSnafu { + reason: format!( + "reclaim boundary {reclaim_boundary} is above the object sequence {object_seq}" + ), + } + ); + let prev = (prev_object_seq != NO_PREDECESSOR).then_some(ChainLink { + object_seq: prev_object_seq, + epoch: prev_epoch, + }); Ok(Header { object_seq, - writer_instance, + epoch, + prev, + reclaim_boundary, }) } @@ -274,12 +313,6 @@ pub(crate) fn decode_footer(bytes: &[u8], trailer: FixedTrailer) -> Result 0, - CorruptedWalObjectSnafu { - reason: "object has no records", - } - ); let expected_len = count .checked_mul(FOOTER_ENTRY_LEN) .and_then(|len| len.checked_add(4)) @@ -551,10 +584,19 @@ pub(crate) fn decode_object(bytes: &[u8]) -> Result { } fn encode_header(header: &Header, output: &mut BytesMut) { + let start = output.len(); output.put_slice(HEADER_MAGIC); output.put_u16(FORMAT_VERSION); output.put_u64(header.object_seq); - output.put_slice(&header.writer_instance); + output.put_u64(header.epoch); + let prev = header.prev.unwrap_or(ChainLink { + object_seq: NO_PREDECESSOR, + epoch: 0, + }); + output.put_u64(prev.object_seq); + output.put_u64(prev.epoch); + output.put_u64(header.reclaim_boundary); + output.put_u32(crc32fast::hash(&output[start..])); } fn encode_footer(directory: &[FooterEntry], output: &mut BytesMut) -> Result<()> { @@ -697,7 +739,12 @@ mod tests { fn header() -> Header { Header { object_seq: 42, - writer_instance: *b"writer-instance!", + epoch: 5, + prev: Some(ChainLink { + object_seq: 40, + epoch: 4, + }), + reclaim_boundary: 37, } } @@ -804,13 +851,36 @@ mod tests { bad_version[8..10].copy_from_slice(&(FORMAT_VERSION + 1).to_be_bytes()); assert_corrupted(decode_object(&bad_version), "unsupported format version 2"); - let mut bad_object_seq = encoded.bytes.to_vec(); - bad_object_seq[10] ^= 1; - assert_corrupted(decode_object(&bad_object_seq), "object checksum mismatch"); + // Every field after the version is covered by the header checksum, + // which recovery verifies without reading the whole object. + for offset in [10, 18, 26, 34, 42, HEADER_LEN - 1] { + let mut bad_field = encoded.bytes.to_vec(); + bad_field[offset] ^= 1; + assert_corrupted(decode_header(&bad_field), "header checksum mismatch"); + } + } - let mut bad_writer = encoded.bytes.to_vec(); - bad_writer[18] ^= 1; - assert_corrupted(decode_object(&bad_writer), "object checksum mismatch"); + #[test] + fn test_format_round_trips_the_first_object_of_a_chain() { + let first = Header { + prev: None, + ..header() + }; + let encoded = encode_object(first.clone(), &records()).unwrap(); + assert_eq!(first, decode_object(&encoded.bytes).unwrap().header); + } + + #[test] + fn test_format_rejects_a_reclaim_boundary_above_the_object() { + let above = Header { + reclaim_boundary: header().object_seq + 1, + ..header() + }; + let encoded = encode_object(above, &records()).unwrap(); + assert_corrupted( + decode_header(&encoded.bytes), + "reclaim boundary 43 is above the object sequence 42", + ); } #[test] @@ -864,41 +934,19 @@ mod tests { } #[test] - fn test_format_rejects_empty_object() { - assert_corrupted(encode_object(header(), &[]), "object has no records"); - } - - #[test] - fn test_format_rejects_checksum_valid_empty_footer() { - let mut object = BytesMut::new(); - encode_header(&header(), &mut object); - let footer_offset = object.len(); - encode_footer(&[], &mut object).unwrap(); - let footer = object[footer_offset..].to_vec(); - let trailer = FixedTrailer { - footer_offset: footer_offset as u64, - footer_len: footer.len() as u64, - footer_crc32: crc32fast::hash(&footer), - object_crc32: 0, - }; - let trailer = FixedTrailer { - object_crc32: object_crc32(&object, trailer), - ..trailer - }; - encode_trailer(trailer, &mut object); - - // Every structural check that recovery runs before the footer passes. - assert_eq!(header(), decode_header(&object).unwrap()); - let trailer_start = object.len() - TRAILER_LEN; - assert_eq!(trailer, decode_trailer(&object[trailer_start..]).unwrap()); + fn test_format_round_trips_an_object_without_records() { + let encoded = encode_object(header(), &[]).unwrap(); + assert!(encoded.footer.is_empty()); assert_eq!( - footer_offset..trailer_start, - footer_range(trailer, object.len()).unwrap() + MIN_OBJECT_LEN, + encoded.bytes.len(), + "an empty object is a header, an empty footer and the trailer" ); - verify_segment_ranges(&[], footer_offset).unwrap(); - - assert_corrupted(decode_footer(&footer, trailer), "object has no records"); - assert_corrupted(decode_object(&object), "object has no records"); + assert_eq!(encoded.bytes.len() as u64, object_len(&encoded.footer)); + let decoded = decode_object(&encoded.bytes).unwrap(); + assert_eq!(header(), decoded.header); + assert!(decoded.records.is_empty()); + assert!(decoded.footer.is_empty()); } #[test] @@ -950,12 +998,17 @@ mod tests { /// to field order, endianness or checksum coverage fails this test even when /// the encoder and the decoder change together. const FIXTURE_V1_HEX: &str = concat!( - // Header: magic, version 1, object sequence 7, writer instance. + // Header: magic, version 1, object sequence 7, epoch 3, predecessor + // sequence 6, predecessor epoch 2, reclaim boundary 5, header CRC32. "475457414c4f424a", "0001", "0000000000000007", - "7772697465722d666978747572652d31", - // Segment of region 1 at offset 34: region id, entry count, then + "0000000000000003", + "0000000000000006", + "0000000000000002", + "0000000000000005", + "118e773f", + // Segment of region 1 at offset 54: region id, entry count, then // (entry id, payload length, payload) per entry. "0000000100000001", "00000003", @@ -968,7 +1021,7 @@ mod tests { "0000000000000003", "00000003", "636363", - // Segment of region 2 at offset 88. + // Segment of region 2 at offset 108. "0000000200000001", "00000002", "000000000000000a", @@ -977,28 +1030,28 @@ mod tests { "000000000000000b", "00000002", "7979", - // Footer at offset 127: entry count, then per segment region id, min and + // Footer at offset 147: entry count, then per segment region id, min and // max entry id, entry count, segment offset, length and CRC32. "00000002", "0000000100000001", "0000000000000001", "0000000000000003", "00000003", - "0000000000000022", + "0000000000000036", "0000000000000036", "25ac0486", "0000000200000001", "000000000000000a", "000000000000000b", "00000002", - "0000000000000058", + "000000000000006c", "0000000000000027", "65dc08ec", // Trailer: footer offset, footer length, footer CRC32, object CRC32, magic. - "000000000000007f", + "0000000000000093", "0000000000000064", - "8298e0f7", - "7a009db1", + "d80c9f07", + "65c91839", "475457414c54524c", ); @@ -1031,7 +1084,12 @@ mod tests { let fixture = fixture_bytes(); let header = Header { object_seq: 7, - writer_instance: *b"writer-fixture-1", + epoch: 3, + prev: Some(ChainLink { + object_seq: 6, + epoch: 2, + }), + reclaim_boundary: 5, }; let footer = vec![ FooterEntry { @@ -1039,7 +1097,7 @@ mod tests { min_entry_id: 1, max_entry_id: 3, entry_count: 3, - segment_offset: 34, + segment_offset: 54, segment_len: 54, segment_crc32: 0x25ac0486, }, @@ -1048,12 +1106,12 @@ mod tests { min_entry_id: 10, max_entry_id: 11, entry_count: 2, - segment_offset: 88, + segment_offset: 108, segment_len: 39, segment_crc32: 0x65dc08ec, }, ]; - assert_eq!(259, fixture.len()); + assert_eq!(279, fixture.len()); let decoded = decode_object(&fixture).unwrap(); assert_eq!(header, decoded.header); @@ -1061,10 +1119,10 @@ mod tests { assert_eq!(fixture_records(), decoded.records); assert_eq!( FixedTrailer { - footer_offset: 127, + footer_offset: 147, footer_len: 100, - footer_crc32: 0x8298e0f7, - object_crc32: 0x7a009db1, + footer_crc32: 0xd80c9f07, + object_crc32: 0x65c91839, }, decode_trailer(&fixture[fixture.len() - TRAILER_LEN..]).unwrap() ); diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs index 0af6ac70278..4fc7e971535 100644 --- a/src/log-store/src/object_store_wal/io.rs +++ b/src/log-store/src/object_store_wal/io.rs @@ -138,6 +138,16 @@ impl ObjectStoreIo { }) } + /// Deletes the object `object_seq`. Deleting an object that does not + /// exist succeeds. + pub(crate) async fn delete(&self, object_seq: u64) -> Result<()> { + let path = self.object_path(object_seq); + self.store.delete(&path).await.context(WalObjectStoreSnafu { + operation: "delete", + path, + }) + } + /// Lists the objects under the prefix, ordered by object sequence. Keys /// that do not follow the object layout are ignored. pub(crate) async fn list(&self) -> Result> { @@ -357,6 +367,11 @@ mod tests { vec![1, 1, 2], objects.iter().map(|object| object.size).collect::>() ); + + // A delete of an object that is already gone succeeds. + io.delete(2).await.unwrap(); + io.delete(2).await.unwrap(); + assert_eq!(vec![1, 10], object_seqs(io.list().await.unwrap())); } #[tokio::test] diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index c0d9cc89363..d3a3acc6968 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -12,18 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Object store WAL construction, recovery and region reads. +//! Object store WAL construction, recovery, writes and region reads. -use std::collections::{BTreeSet, HashMap}; +use std::cmp::Reverse; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; use std::fmt; use std::ops::Range; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, PoisonError, RwLock}; -use std::time::Duration; +use std::time::{Duration, Instant}; use async_stream::try_stream; use bytes::Bytes; -use common_wal::config::object_store::ObjectStoreWalConfig; +use common_telemetry::{debug, info, warn}; +use common_wal::config::object_store::{AckMode, CorruptedSegmentAction, ObjectStoreWalConfig}; +use futures::future::BoxFuture; +use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt}; use object_store::ObjectStore; use snafu::{IntoError, OptionExt, ResultExt, ensure}; @@ -31,22 +35,52 @@ use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::{ObjectStoreProvider, Provider}; use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; use store_api::storage::RegionId; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{mpsc, oneshot, watch}; +use tokio::time::MissedTickBehavior; use crate::error::{ - CorruptedWalObjectSnafu, Error, InvalidProviderSnafu, InvalidWalObjectSnafu, - InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, MismatchedWalRegionSnafu, - ObjectStoreWalSnafu, Result, UnsupportedObjectStoreWalOperationSnafu, + CorruptedWalObjectSnafu, Error, IncompleteWalEntrySnafu, InvalidProviderSnafu, + InvalidWalObjectSnafu, InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, + MismatchedWalRegionSnafu, ObjectStoreWalSnafu, ObjectStoreWalStoppedSnafu, Result, + StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, + WalObjectSequenceUnsettledSnafu, }; +use crate::metrics::{ + METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS, METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL, + METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL, METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL, + METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL, + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES, METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS, + METRIC_OBJECT_STORE_WAL_OBJECT_BYTES, METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES, + METRIC_OBJECT_STORE_WAL_POISONED_TOTAL, METRIC_OBJECT_STORE_WAL_READ_SECONDS, + METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS, + METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL, METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS, + METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL, +}; +use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, entry_id, sequence_floor}; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ - FixedTrailer, FooterEntry, HEADER_LEN, MIN_OBJECT_LEN, TRAILER_LEN, decode_footer, - decode_header, decode_segment, decode_trailer, footer_range, verify_segment_ranges, + ChainLink, EncodedObject, FixedTrailer, FooterEntry, HEADER_LEN, Header, MIN_OBJECT_LEN, + Record, TRAILER_LEN, decode_footer, decode_header, decode_segment, decode_trailer, + encode_object, footer_range, object_len, verify_segment_ranges, }; use crate::object_store_wal::io::{ListedObject, ObjectStoreIo, PutResult}; const COMMAND_BUFFER: usize = 1024; +/// Appends queued for the actor; a caller beyond them waits to send. +const APPEND_BUFFER: usize = 16; const MIN_FLUSH_INTERVAL: Duration = Duration::from_millis(10); +const MAX_IN_FLIGHT_CREATES: usize = 4; +const MAX_IN_FLIGHT_DELETES: usize = 4; +/// Delay before a create that failed transiently is attempted again in the +/// `enqueued` acknowledgement mode, where no caller is left to retry it. +const CREATE_RETRY_DELAY: Duration = Duration::from_millis(100); +/// Number of sealed batches that may wait to be created or indexed. One +/// admission can seal two batches, the open batch before an append that would +/// exhaust its positions and the append's own batch, so the actor takes an +/// append from its channel only while two more fit: a stalled object store +/// holds callers back instead of growing the backlog. +const MAX_SEALED_BATCHES: usize = 2 * MAX_IN_FLIGHT_CREATES; /// Number of objects whose footers recovery fetches at a time. const RECOVERY_CONCURRENCY: usize = 8; /// Bytes recovery reads from the end of an object in one request. The window @@ -55,33 +89,148 @@ const RECOVERY_CONCURRENCY: usize = 8; const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// A log store over the immutable WAL objects under one prefix. -pub(crate) struct ObjectStoreLogStore { +/// +/// Appends are admitted into an open batch. A background actor seals the batch +/// when it reaches the size limit or the flush interval elapses and creates +/// the object under the next sequence while it keeps admitting entries into +/// the next batch; up to [`MAX_IN_FLIGHT_CREATES`] creates run at a time. +/// Objects are indexed in the catalog in sequence order. In the `durable` +/// acknowledgement mode an append returns once its object is durable and +/// indexed, so an acknowledged entry never has a missing predecessor; in the +/// `enqueued` mode it returns on admission and the object is created in the +/// background. While [`MAX_SEALED_BATCHES`] batches +/// wait to become durable no append is admitted. A segment that a read cannot +/// decode is handled as `on_corrupted_segment` says. Objects below the reclaim +/// boundary of the last indexed object are deleted in the background. +pub struct ObjectStoreLogStore { prefix: String, + ack_mode: AckMode, + on_corrupted_segment: CorruptedSegmentAction, io: Arc, catalog: Arc>, obsolete_entry_ids: ObsoleteEntryIds, + /// Segments that reads skipped, per region and by object sequence. + wal_holes: WalHoles, + /// Objects with a delete in flight, which a read consults. + deleting: Arc, terminal_error: TerminalError, stopped: Arc, command_tx: mpsc::Sender, + append_tx: mpsc::Sender, + #[cfg(any(test, feature = "testing"))] + admitted_appends: watch::Receiver, + #[cfg(any(test, feature = "testing"))] + creates_held: watch::Sender, + #[cfg(any(test, feature = "testing"))] + creates_fail: Arc, + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write: Arc, + #[cfg(any(test, feature = "testing"))] + parked_creates: watch::Receiver, + #[cfg(any(test, feature = "testing"))] + durability_waits: watch::Receiver, } type ObsoleteEntryIds = Arc>>; +type WalHoles = Arc>>>; + +/// A segment a read skipped because it did not decode: the entries of one +/// region in one object are missing from every read of that region. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WalHole { + /// Key of the object that holds the segment. + pub path: String, + pub object_seq: u64, + /// Entry id range the footer records for the segment. + pub min_entry_id: EntryId, + pub max_entry_id: EntryId, +} + type TerminalError = Arc>>>; +/// The delete attempts in flight, one per object. +/// +/// The actor registers an attempt before its delete starts, unindexes the +/// object if the delete succeeded, and only then settles the attempt and takes +/// it out, so an object without an attempt was either never collected or is +/// already unindexed. [`is_collected`] depends on that order. +#[derive(Debug, Default)] +struct DeletingObjects { + /// The settlement flag of the attempt in flight per object sequence. + attempts: Mutex>>, +} + +impl DeletingObjects { + fn lock(&self) -> std::sync::MutexGuard<'_, BTreeMap>> { + self.attempts.lock().unwrap_or_else(PoisonError::into_inner) + } + + fn start(&self, object_seq: u64) { + self.lock().insert(object_seq, watch::channel(false).0); + } + + fn contains(&self, object_seq: u64) -> bool { + self.lock().contains_key(&object_seq) + } + + /// Settles the attempt of `object_seq` and takes it out, releasing the + /// reads that wait for it. + fn finish(&self, object_seq: u64) { + if let Some(settled) = self.lock().remove(&object_seq) { + settled.send_replace(true); + } + } + + /// Settles every attempt the actor will never complete because it exited: + /// their objects stay indexed, so waiting reads report their error. + fn abandon_all(&self) { + for (_, settled) in std::mem::take(&mut *self.lock()) { + settled.send_replace(true); + } + } + + /// Returns the settlement of the attempt in flight for `object_seq`, or + /// `None` when no delete of it is in flight. + fn attempt(&self, object_seq: u64) -> Option> { + self.lock().get(&object_seq).map(watch::Sender::subscribe) + } + + /// Waits until no delete is in flight. + #[cfg(any(test, feature = "testing"))] + async fn wait_until_empty(&self) { + loop { + let in_flight = self + .lock() + .values() + .map(watch::Sender::subscribe) + .collect::>(); + if in_flight.is_empty() { + return; + } + for mut settled in in_flight { + let _ = settled.wait_for(|settled| *settled).await; + } + } + } +} + impl fmt::Debug for ObjectStoreLogStore { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ObjectStoreLogStore") .field("prefix", &self.prefix) + .field("ack_mode", &self.ack_mode) + .field("on_corrupted_segment", &self.on_corrupted_segment) .finish_non_exhaustive() } } impl ObjectStoreLogStore { /// Builds the store under the node and generation prefix derived from `config`, - /// recovering the catalog from the objects that already exist. Recovery - /// fails on the first corrupted or conflicting object. - pub(crate) async fn try_new( + /// recovering the catalog from the objects that already exist and writing + /// the object that starts the epoch of this instance. Recovery fails on the + /// first corrupted or conflicting object. + pub async fn try_new( object_store: ObjectStore, config: &ObjectStoreWalConfig, node_id: u64, @@ -107,37 +256,192 @@ impl ObjectStoreLogStore { } ); - positive_bytes(config.max_batch_bytes.as_bytes(), "max batch bytes")?; - let (catalog, next_object_seq, durable_entry_ids) = recover(io.as_ref()).await?; + let max_batch_bytes = positive_bytes(config.max_batch_bytes.as_bytes(), "max batch bytes")?; + let max_unpersisted_bytes = positive_bytes( + config.max_unpersisted_bytes.as_bytes(), + "max unpersisted bytes", + )?; + ensure!( + config.max_unpersisted_age > Duration::ZERO, + InvalidWalObjectStoreSnafu { + reason: "max unpersisted age is zero", + } + ); + let recovery_started_at = Instant::now(); + let Recovered { + mut catalog, + next_object_seq, + durable_entry_ids, + tip, + max_epoch, + reclaim_boundary, + unindexed, + recovered_objects, + } = recover(io.as_ref()).await?; + // Every epoch is one above the sequence of the start object its + // instance created, and every object is at or above its start object, + // so an epoch above the next sequence names no instance that ran. + ensure!( + max_epoch <= next_object_seq, + CorruptedWalObjectSnafu { + reason: format!( + "an object carries epoch {max_epoch}, above the next sequence {next_object_seq}" + ), + } + ); + // Only a recovery that succeeded is reported, so that the objects + // counted and the time measured are those of the same recoveries. + METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL.inc_by(recovered_objects); + METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS + .observe(recovery_started_at.elapsed().as_secs_f64()); + publish_indexed(&catalog); + let (start, moved_past) = + start_epoch(io.as_ref(), next_object_seq, tip, reclaim_boundary).await?; + let epoch = start.epoch; + // The epoch identifies this instance in every object it writes. + info!( + "Opened object store WAL under {prefix} at epoch {epoch}, start object {}", + start.object_seq + ); + catalog + .insert_object(start.object_seq, Vec::new()) + .with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(start.object_seq), + })?; + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.inc(); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.add(object_len(&[]) as i64); + // No watermark is recorded yet. + let writer_reclaim_boundary = catalog.reclaim_boundary(reclaim_boundary, &HashMap::new()); let catalog = Arc::new(RwLock::new(catalog)); + let obsolete_entry_ids = ObsoleteEntryIds::default(); + let wal_holes = WalHoles::default(); + let deleting = Arc::new(DeletingObjects::default()); let terminal_error = TerminalError::default(); let stopped = Arc::new(AtomicBool::new(false)); let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); + let (append_tx, append_rx) = mpsc::channel(APPEND_BUFFER); + #[cfg(any(test, feature = "testing"))] + let (admitted_appends_tx, admitted_appends_rx) = watch::channel(0); + #[cfg(any(test, feature = "testing"))] + let (creates_held_tx, creates_held_rx) = watch::channel(false); + #[cfg(any(test, feature = "testing"))] + let creates_fail = Arc::new(AtomicBool::new(false)); + #[cfg(any(test, feature = "testing"))] + let next_create_fails_after_write = Arc::new(AtomicBool::new(false)); + #[cfg(any(test, feature = "testing"))] + let (parked_creates_tx, parked_creates_rx) = watch::channel(0); + #[cfg(any(test, feature = "testing"))] + let (durability_waits_tx, durability_waits_rx) = watch::channel(0); + let actor = Actor { + io: io.clone(), + catalog: catalog.clone(), + obsolete_entry_ids: obsolete_entry_ids.clone(), + wal_holes: wal_holes.clone(), terminal_error: terminal_error.clone(), stopped: stopped.clone(), command_rx, + append_rx, + ack_mode: config.ack_mode, + max_unpersisted_bytes, + max_unpersisted_age: config.max_unpersisted_age, + open_batch: OpenBatch::new(max_batch_bytes), issued_entry_ids: durable_entry_ids, - next_object_seq: Some(next_object_seq), + pending: Vec::new(), + sealed: VecDeque::new(), + creates: FuturesUnordered::new(), + stalled: None, + durable_waiters: Vec::new(), stop: Vec::new(), + stop_error: None, + next_object_seq: start + .object_seq + .checked_add(1) + .filter(|next_object_seq| *next_object_seq < OBJECT_SEQ_LIMIT), + epoch, + last_indexed: start, + reclaim_boundary: writer_reclaim_boundary, + durable_reclaim_boundary: reclaim_boundary, + unindexed: unindexed.into_iter().chain(moved_past).collect(), + creating: HashSet::new(), + deleting: deleting.clone(), + deletes: FuturesUnordered::new(), + failed_deletes: HashSet::new(), + flush_interval: config.flush_interval, + #[cfg(any(test, feature = "testing"))] + admitted_appends: admitted_appends_tx, + #[cfg(any(test, feature = "testing"))] + creates_held: creates_held_rx, + #[cfg(any(test, feature = "testing"))] + creates_fail: creates_fail.clone(), + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write: next_create_fails_after_write.clone(), + #[cfg(any(test, feature = "testing"))] + parked_creates: Arc::new(parked_creates_tx), + #[cfg(any(test, feature = "testing"))] + durability_waits: durability_waits_tx, }; + #[cfg(test)] + LIVE_ACTORS.fetch_add(1, Ordering::SeqCst); common_runtime::spawn_global(actor.run()); Ok(Arc::new(Self { prefix, + ack_mode: config.ack_mode, + on_corrupted_segment: config.on_corrupted_segment, io, catalog, - obsolete_entry_ids: ObsoleteEntryIds::default(), + obsolete_entry_ids, + wal_holes, + deleting, terminal_error, stopped, command_tx, + append_tx, + #[cfg(any(test, feature = "testing"))] + admitted_appends: admitted_appends_rx, + #[cfg(any(test, feature = "testing"))] + creates_held: creates_held_tx, + #[cfg(any(test, feature = "testing"))] + creates_fail, + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write, + #[cfg(any(test, feature = "testing"))] + parked_creates: parked_creates_rx, + #[cfg(any(test, feature = "testing"))] + durability_waits: durability_waits_rx, })) } + /// Returns the largest entry id of the provider's region whose object is + /// durable and indexed, or zero for a region without such entries. In the + /// `enqueued` acknowledgement mode an entry id that an append returned + /// stays above this value until the object holding it is created. + pub(crate) fn durable_entry_id(&self, provider: &Provider) -> Result { + self.check_terminal()?; + let region_id = self.region_of(provider)?; + Ok(self.durable_entry_id_of(region_id)) + } + fn durable_entry_id_of(&self, region_id: RegionId) -> EntryId { let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); catalog.region_max_entry_id(region_id).unwrap_or(0) } + /// Returns the segments of the provider's region that reads of this store + /// skipped because they did not decode, ordered by object sequence. Holes + /// are kept in memory only: a store opened later on the same prefix + /// records them again when its reads meet the segments. + pub fn wal_holes(&self, provider: &Provider) -> Result> { + let region_id = self.region_of(provider)?; + Ok(self + .wal_holes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(®ion_id) + .map(|holes| holes.values().cloned().collect()) + .unwrap_or_default()) + } + /// Returns the region of `provider`, which must select this store's prefix. fn region_of(&self, provider: &Provider) -> Result { let provider = @@ -178,6 +482,20 @@ impl ObjectStoreLogStore { } } +/// The transient object store error of a create that a testing hook fails. +#[cfg(any(test, feature = "testing"))] +fn injected_create_failure(io: &dyn WalObjectIo, object_seq: u64) -> Result { + let error = object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected create failure", + ) + .set_temporary(); + Err(error).context(crate::error::WalObjectStoreSnafu { + operation: "write", + path: io.object_path(object_seq), + }) +} + /// Records the obsolete watermark of `region_id`, which never moves down. fn record_obsolete(obsolete_entry_ids: &ObsoleteEntryIds, region_id: RegionId, entry_id: EntryId) { obsolete_entry_ids @@ -197,10 +515,154 @@ fn positive_bytes(bytes: u64, name: &str) -> Result { }) } +#[cfg(any(test, feature = "testing"))] +impl ObjectStoreLogStore { + /// Waits until the actor has admitted at least `expected` append calls + /// since the store was built. + pub async fn wait_for_admitted_appends(&self, expected: usize) -> Result<()> { + self.admitted_appends + .clone() + .wait_for(|count| *count >= expected) + .await + .ok() + .map(|_| ()) + .context(ObjectStoreWalStoppedSnafu) + } + + /// Seals the open batch regardless of its size and age and returns once + /// its object is durable and indexed, or with the error that failed it. + pub async fn seal_open_batch(&self) -> Result<()> { + ensure!( + !self.stopped.load(Ordering::Acquire), + ObjectStoreWalStoppedSnafu + ); + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(Command::Seal { + response: response_tx, + }) + .await + .ok() + .context(ObjectStoreWalStoppedSnafu)?; + response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? + } + + /// Parks every conditional create that starts from now on until + /// [`release_creates`](Self::release_creates), so a test can observe + /// entries that are admitted but not durable. A create that is parked when + /// the store is dropped never runs. + pub fn hold_creates(&self) { + self.creates_held.send_replace(true); + } + + /// Lets the creates parked by [`hold_creates`](Self::hold_creates) run. + pub fn release_creates(&self) { + self.creates_held.send_replace(false); + } + + /// Makes every create that runs from now on fail with a transient object + /// store error instead of writing. + pub fn fail_creates(&self) { + self.creates_fail.store(true, Ordering::Release); + } + + /// Makes the next create that writes its object report a transient object + /// store error afterwards, so the object exists although its create + /// failed. + pub fn fail_next_create_after_write(&self) { + self.next_create_fails_after_write + .store(true, Ordering::Release); + } + + /// Waits until at least `expected` creates have been parked by + /// [`hold_creates`](Self::hold_creates) since the store was built. + pub async fn wait_for_parked_creates(&self, expected: usize) -> Result<()> { + self.parked_creates + .clone() + .wait_for(|count| *count >= expected) + .await + .ok() + .map(|_| ()) + .context(ObjectStoreWalStoppedSnafu) + } + + /// Waits until at least `expected` calls of + /// [`wait_durable`](LogStore::wait_durable) have had to wait for an entry + /// that is not durable since the store was built. + pub async fn wait_for_durability_waits(&self, expected: usize) -> Result<()> { + self.durability_waits + .clone() + .wait_for(|count| *count >= expected) + .await + .ok() + .map(|_| ()) + .context(ObjectStoreWalStoppedSnafu) + } + + /// Ends the actor the way a crash of the process would: nothing more is + /// written, creates that have not completed are dropped, and every caller + /// still waiting for the store fails. Returns once the actor has exited. + pub async fn crash(&self) { + self.stopped.store(true, Ordering::Release); + let (response_tx, response_rx) = oneshot::channel(); + if self + .command_tx + .send(Command::Crash { + response: response_tx, + }) + .await + .is_ok() + { + let _ = response_rx.await; + } + } + + /// Waits until no object delete is in flight. The deletes a collection + /// starts are in flight by the time the `obsolete`, `obsolete_all`, + /// [`seal_open_batch`](Self::seal_open_batch) or `durable` append that + /// started it returns. + pub async fn wait_for_garbage_collection(&self) { + self.deleting.wait_until_empty().await; + } + + /// Sets the stopped flag without sending the stop command, which is the + /// state a store is in between the two steps of [`stop`](LogStore::stop). + pub fn begin_stop(&self) { + self.stopped.store(true, Ordering::Release); + } + + /// Returns the key of the object `object_seq` and the byte range the + /// segment of the provider's region occupies in it, so a test can damage + /// that segment alone, or `None` when the object holds no segment of the + /// region. + pub fn segment_location( + &self, + provider: &Provider, + object_seq: u64, + ) -> Result)>> { + let region_id = self.region_of(provider)?; + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + Ok(catalog + .objects_for_entry_range(region_id, 0, EntryId::MAX)? + .into_iter() + .find(|(seq, _)| *seq == object_seq) + .map(|(_, entry)| { + let start = entry.segment_offset; + ( + self.io.object_path(object_seq), + start..start + entry.segment_len, + ) + })) + } +} + #[async_trait::async_trait] impl LogStore for ObjectStoreLogStore { type Error = Error; + /// Stops the store. Creates in flight run to completion and acknowledge + /// their entries if they succeed; in the `enqueued` mode the remaining + /// backlog is uploaded first, and a failure of that upload is returned. async fn stop(&self) -> Result<()> { self.stopped.store(true, Ordering::Release); let (response_tx, response_rx) = oneshot::channel(); @@ -217,12 +679,49 @@ impl LogStore for ObjectStoreLogStore { Ok(()) } - async fn append_batch(&self, _entries: Vec) -> Result { - UnsupportedObjectStoreWalOperationSnafu.fail() + async fn append_batch(&self, entries: Vec) -> Result { + ensure!( + !self.stopped.load(Ordering::Acquire), + ObjectStoreWalStoppedSnafu + ); + self.check_terminal()?; + if entries.is_empty() { + return Ok(AppendBatchResponse::default()); + } + for entry in &entries { + let region_id = self.region_of(entry.provider())?; + ensure!( + region_id == entry.region_id(), + MismatchedWalRegionSnafu { + region_id: entry.region_id(), + reason: format!("provider belongs to region {region_id}"), + } + ); + ensure!(entry.is_complete(), IncompleteWalEntrySnafu { region_id }); + } + + let requested_at = Instant::now(); + let (response_tx, response_rx) = oneshot::channel(); + self.append_tx + .send((entries, response_tx)) + .await + .ok() + .context(ObjectStoreWalStoppedSnafu)?; + let response = response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)?; + // Only an acknowledgement is timed; an append that failed never got one. + if self.ack_mode == AckMode::Durable && response.is_ok() { + METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS + .observe(requested_at.elapsed().as_secs_f64()); + } + response } /// Reads the provider's region from `entry_id`, hiding obsolete entries. /// The catalog locates each segment, so a caller-supplied index is unnecessary. + /// An object collected after the read listed it is skipped. Otherwise a + /// segment that does not decode fails the read or is skipped and recorded + /// as a hole, as `on_corrupted_segment` says; a failed fetch always fails + /// the read. async fn read( &self, provider: &Provider, @@ -256,16 +755,32 @@ impl LogStore for ObjectStoreLogStore { }; let io = self.io.clone(); + let catalog = self.catalog.clone(); + let deleting = self.deleting.clone(); let provider = provider.clone(); + let on_corrupted_segment = self.on_corrupted_segment; + let wal_holes = self.wal_holes.clone(); Ok(Box::pin(try_stream! { + let started_at = Instant::now(); for (object_seq, footer_entry) in objects { - let bytes = io - .get_range(object_seq, footer_entry.segment_offset, footer_entry.segment_len) - .await?; - let records = decode_segment(&bytes, &footer_entry) - .with_context(|_| InvalidWalObjectSnafu { - path: io.object_path(object_seq), - })?; + let records = match fetch_segment(io.as_ref(), object_seq, &footer_entry).await { + Ok(records) => records, + Err(error) if is_collected(&catalog, &deleting, object_seq).await => { + debug!( + "Skipped WAL object {} (sequence {}), collected after the read listed it: {error}", + io.object_path(object_seq), + object_seq + ); + continue; + } + Err(error @ Error::InvalidWalObject { .. }) + if on_corrupted_segment == CorruptedSegmentAction::Skip => + { + record_skipped_segment(&catalog, &wal_holes, io.as_ref(), object_seq, &footer_entry, &error); + continue; + } + Err(error) => Err(error)?, + }; let entries = records .into_iter() .filter(|record| record.entry_id >= start_entry_id) @@ -282,6 +797,7 @@ impl LogStore for ObjectStoreLogStore { yield entries; } } + METRIC_OBJECT_STORE_WAL_READ_SECONDS.observe(started_at.elapsed().as_secs_f64()); })) } @@ -308,7 +824,26 @@ impl LogStore for ObjectStoreLogStore { .collect()) } - /// Records a memory-only watermark. The caller re-establishes it after a restart. + /// Moves the obsolete watermark of the region up to `entry_id`. In the + /// `enqueued` acknowledgement mode the watermark never passes the durable + /// entry id: an entry that is not durable yet is replayed after a crash, + /// and hiding it would skip it. + /// + /// Every id the region is assigned from now on is greater than + /// `entry_id`, whatever watermark is recorded: the sequence of the next + /// object is raised above the object `entry_id` names unless it is there + /// already, so a watermark assigned under another prefix, or one whose + /// object this prefix no longer holds, is never passed by a new id. Zero + /// names no object and needs no floor. + /// + /// The watermark and the floor are applied together by the actor: a call + /// cancelled before its command is queued publishes neither, while a + /// command already queued still applies both. While the open batch holds + /// ids under the next sequence the floor cannot move, and the call fails + /// with [`Error::WalObjectSequenceUnsettled`] and records neither. The + /// watermark is kept in memory; the caller re-establishes it after a + /// restart. In the `enqueued` mode an `entry_id` this store handed out + /// needs no floor: it is never handed out again while the store runs. async fn obsolete( &self, provider: &Provider, @@ -316,15 +851,49 @@ impl LogStore for ObjectStoreLogStore { entry_id: EntryId, ) -> Result<()> { self.check_region(provider, region_id)?; - record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); - Ok(()) + let watermark = match self.ack_mode { + AckMode::Durable => entry_id, + AckMode::Enqueued => entry_id.min(self.durable_entry_id_of(region_id)), + }; + let (response_tx, response_rx) = oneshot::channel(); + let command = Command::Obsolete { + region_id, + entry_id, + watermark, + response: response_tx, + }; + let answered = match self.command_tx.send(command).await { + Ok(()) => response_rx.await.ok(), + Err(_) => None, + }; + match answered { + Some(result) => result, + // The actor exited, before the command was queued or with it + // still queued behind the stop: nothing is assigned an id any + // more, so the watermark alone is consistent. + None => { + record_obsolete(&self.obsolete_entry_ids, region_id, watermark); + Ok(()) + } + } } - /// Hides every entry of the region in memory. The caller re-establishes - /// the watermark after a restart. + /// Hides every entry of the region in memory and lets collection delete + /// its objects. The caller re-establishes the watermark after a restart. async fn obsolete_all(&self, provider: &Provider, region_id: RegionId) -> Result<()> { self.check_region(provider, region_id)?; record_obsolete(&self.obsolete_entry_ids, region_id, EntryId::MAX); + let (response_tx, response_rx) = oneshot::channel(); + if self + .command_tx + .send(Command::Collect { + response: response_tx, + }) + .await + .is_ok() + { + let _ = response_rx.await; + } Ok(()) } @@ -345,56 +914,1124 @@ impl LogStore for ObjectStoreLogStore { } fn latest_entry_id(&self, provider: &Provider) -> Result { + self.durable_entry_id(provider) + } + + /// Waits until every entry of the provider's region with an id at or + /// below `entry_id` is durable and indexed. Returns at once in the + /// `durable` acknowledgement mode, where a caller only holds durable ids. + async fn wait_durable(&self, provider: &Provider, entry_id: EntryId) -> Result<()> { self.check_terminal()?; let region_id = self.region_of(provider)?; - Ok(self.durable_entry_id_of(region_id)) + if entry_id <= self.durable_entry_id_of(region_id) { + return Ok(()); + } + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(Command::WaitDurable { + region_id, + entry_id, + response: response_tx, + }) + .await + .ok() + .context(ObjectStoreWalStoppedSnafu)?; + response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? } } +type AppendResponse = oneshot::Sender>; + +type QueuedAppend = (Vec, AppendResponse); + enum Command { + /// Answered once the region is durable through `entry_id`. + WaitDurable { + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, + }, + /// Answered once the watermark is recorded and no id of the region at + /// or below `entry_id` can be assigned, or with the reason neither was + /// done. + Obsolete { + region_id: RegionId, + entry_id: EntryId, + watermark: EntryId, + response: oneshot::Sender>, + }, + /// Answered once the collection the recorded watermarks allow started. + Collect { response: oneshot::Sender<()> }, Stop { response: oneshot::Sender>, }, + #[cfg(any(test, feature = "testing"))] + Seal { + response: oneshot::Sender>, + }, + #[cfg(any(test, feature = "testing"))] + Crash { response: oneshot::Sender<()> }, } +/// An append waiting for the object that holds its entries. +struct PendingAppend { + last_entry_ids: HashMap, + response: AppendResponse, +} + +/// A caller waiting until a region is durable through an entry id. +struct DurableWaiter { + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, +} + +/// Where the conditional create of a sealed batch stands. +enum CreateState { + /// The create has not started: no slot was free, or in the `enqueued` + /// mode an earlier attempt failed transiently and is repeated. + Pending, + InFlight, + Created, +} + +/// A sealed and encoded batch that holds its object sequence and waits to be +/// created, indexed and acknowledged. +struct SealedBatch { + object_seq: u64, + reclaim_boundary: u64, + bytes: Bytes, + footer: Vec, + first_admitted_at: Instant, + sealed_at: Instant, + /// Number of creates that were attempted for the batch. + attempts: u32, + waiters: Vec, + #[cfg(any(test, feature = "testing"))] + seal_waiters: Vec>>, + state: CreateState, +} + +impl SealedBatch { + fn is_in_flight(&self) -> bool { + matches!(self.state, CreateState::InFlight) + } + + fn fail(self, error: impl Fn() -> Error) { + for waiter in self.waiters { + let _ = waiter.response.send(Err(error())); + } + #[cfg(any(test, feature = "testing"))] + for waiter in self.seal_waiters { + let _ = waiter.send(Err(error())); + } + } +} + +type CreateOutcome = (u64, Result); + +/// The actor that owns the open batch and the sealed batches until they are +/// durable. +/// +/// Sealed batches form a pipeline in sequence order. A batch is created under +/// its sequence as soon as one of [`MAX_IN_FLIGHT_CREATES`] slots is free, but +/// it is indexed and acknowledged only once every earlier batch is, so the +/// acknowledged history of a region never has a missing predecessor. Every +/// batch links to its pipeline predecessor, the batch sealed before it, or to +/// the last indexed object when no earlier batch waits, so recovery replays a +/// batch only on the chain of a later object. The outcomes of a create are: +/// +/// | Situation | Sequence | Waiters | Store | +/// | --- | --- | --- | --- | +/// | created, or identical retry | advances | acknowledged in order | healthy | +/// | transient error, `durable` mode | never reused: a create that reported an error may have written its object | every batch that is not indexed fails, created ones included; retries are assigned new ids | healthy: the next batch links to the last indexed object, so no chain reaches a failed batch | +/// | transient error, `enqueued` mode | unchanged | already acknowledged | healthy: the create is repeated with the same bytes after [`CREATE_RETRY_DELAY`], which an identical retry accepts | +/// | transient error, `enqueued` mode after `stop` began | never reused | already acknowledged | the backlog is dropped and `stop` reports the error | +/// | an object of an earlier epoch holds the sequence, `durable` mode | never reused | as for a transient error | healthy: that object can never be on the chain | +/// | an object of the same or a later epoch holds the sequence, any conflicting object in the `enqueued` mode, encoding or catalog error | unchanged | every batch that is not indexed fails | poisoned | +/// | created at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence | +/// +/// A create still in flight when its batch failed runs to completion and its +/// outcome is ignored: whatever it stored is off the chain. After `stop` began +/// nothing is admitted and no create starts, except that the `enqueued` mode +/// uploads its backlog; creates in flight run to completion and acknowledge +/// if they succeed. +/// +/// Appends arrive on their own channel, which the actor does not read while +/// [`MAX_SEALED_BATCHES`] batches wait or an append is held back at a backlog +/// threshold of the `enqueued` mode, so commands such as `stop` are handled +/// while admission is held back. +/// +/// The actor also collects garbage. Every batch carries the reclaim boundary +/// current when it was sealed, see [`ObjectCatalog::reclaim_boundary`], which +/// is recomputed after every `obsolete` and every indexed object and never +/// moves down. The boundary of the last indexed object is durable: every +/// present object below it, indexed or not, is deleted in the background, up +/// to [`MAX_IN_FLIGHT_DELETES`] at a time, after an `obsolete` or an indexing +/// that moved it. An indexed object is unindexed once its delete succeeded; a +/// failed delete is attempted again at the next collection. Nothing is +/// collected once stop began or the store is poisoned, and deletes in flight +/// complete before `stop` returns. struct Actor { + io: Arc, + catalog: Arc>, + obsolete_entry_ids: ObsoleteEntryIds, + wal_holes: WalHoles, terminal_error: TerminalError, stopped: Arc, command_rx: mpsc::Receiver, + append_rx: mpsc::Receiver, + ack_mode: AckMode, + max_unpersisted_bytes: usize, + max_unpersisted_age: Duration, + open_batch: OpenBatch, + /// Largest entry id ever handed out per region, whether it became + /// durable or failed. Unlike the accepted ids of the open batch it never + /// moves down. issued_entry_ids: HashMap, - next_object_seq: Option, + /// Waiters of the open batch in the `durable` mode. + pending: Vec, + /// Batches that are not durable yet, in sequence order. + sealed: VecDeque, + /// Creates in flight, including those of batches that already failed. + creates: FuturesUnordered>, + /// The append held back in the `enqueued` mode while the unpersisted + /// backlog is at a threshold, and when it was held back. Later appends + /// wait in their channel. + stalled: Option<(QueuedAppend, Instant)>, + durable_waiters: Vec, + /// Callers of `stop`, answered once nothing is in flight. stop: Vec>>, + /// The failure `stop` reports in the `enqueued` mode once an + /// acknowledged backlog was dropped, recorded when it happens. + stop_error: Option>, + /// Sequence of the next sealed batch, `None` once the sequence is + /// exhausted. The open batch assigns its entry ids from it. It is above + /// every sealed batch and never moves back while the store runs. + next_object_seq: Option, + /// Epoch of this instance, carried by every object it writes. + epoch: u64, + /// The last indexed object, which is the start object until a batch is + /// indexed. A batch sealed while no earlier batch waits extends it. + last_indexed: ChainLink, + /// The reclaim boundary the next sealed batch carries. + reclaim_boundary: u64, + /// The reclaim boundary of the last indexed object. + durable_reclaim_boundary: u64, + /// Sequences of objects that may be present but are not indexed: the + /// objects recovery listed off the chain and the batches that failed. + unindexed: BTreeSet, + /// Sequences whose create is in flight. + creating: HashSet, + deleting: Arc, + deletes: FuturesUnordered)>>, + /// Deletes that failed since the last collection started. + failed_deletes: HashSet, + flush_interval: Duration, + #[cfg(any(test, feature = "testing"))] + admitted_appends: watch::Sender, + #[cfg(any(test, feature = "testing"))] + creates_held: watch::Receiver, + #[cfg(any(test, feature = "testing"))] + creates_fail: Arc, + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write: Arc, + #[cfg(any(test, feature = "testing"))] + parked_creates: Arc>, + #[cfg(any(test, feature = "testing"))] + durability_waits: watch::Sender, } impl Actor { async fn run(mut self) { - while let Some(Command::Stop { response }) = self.command_rx.recv().await { - self.handle_stop(response); + let mut interval = tokio::time::interval(self.flush_interval); + interval.set_missed_tick_behavior(MissedTickBehavior::Delay); + // The first tick completes immediately. + interval.tick().await; + + loop { + tokio::select! { + _ = interval.tick() => { + // Nothing starts after stop began; the `enqueued` mode + // sealed its backlog when stop was requested. + if !self.is_stopped() { + self.flush_open_batch(); + } + } + Some((object_seq, result)) = self.creates.next(), if !self.creates.is_empty() => { + self.on_create_completed(object_seq, result); + } + Some((object_seq, result)) = self.deletes.next(), if !self.deletes.is_empty() => { + self.on_delete_completed(object_seq, result); + } + Some((entries, response)) = self.append_rx.recv(), if self.stalled.is_none() && self.sealed.len() + 2 <= MAX_SEALED_BATCHES => { + self.handle_append(entries, response); + } + command = self.command_rx.recv() => match command { + Some(Command::WaitDurable { region_id, entry_id, response }) => { + self.handle_wait_durable(region_id, entry_id, response); + } + Some(Command::Obsolete { region_id, entry_id, watermark, response }) => { + self.handle_obsolete(region_id, entry_id, watermark, response); + } + Some(Command::Collect { response }) => { + self.update_reclaim_boundary(); + self.collect(); + let _ = response.send(()); + } + Some(Command::Stop { response }) => { + self.handle_stop(response); + } + #[cfg(any(test, feature = "testing"))] + Some(Command::Seal { response }) => { + self.handle_seal(response); + } + #[cfg(any(test, feature = "testing"))] + Some(Command::Crash { response }) => { + self.handle_crash(); + let _ = response.send(()); + return; + } + // Every sender is gone: the store was dropped without + // `stop`. The creates in flight are dropped with the actor. + None => return, + }, + } if self.finish_stop() { return; } } } - fn handle_stop(&mut self, response: oneshot::Sender>) { - self.stop.push(response); + fn handle_append(&mut self, entries: Vec, response: AppendResponse) { + // The append was queued before `stop` set the flag; nothing that is + // not durable yet gets admitted once it is set. + if self.is_stopped() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + return; + } + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + return; + } + if self.ack_mode == AckMode::Enqueued && self.backlog_at_threshold() { + METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.inc(); + self.stalled = Some(((entries, response), Instant::now())); + self.ensure_create_in_flight(); + return; + } + self.admit(entries, response); } - fn finish_stop(&mut self) -> bool { - if !self.is_stopped() || self.stop.is_empty() { + /// Admits `entries` into the open batch, which assigns their ids under + /// the next object sequence. In the `durable` mode the caller waits for + /// the object, in the `enqueued` mode it is answered now. + fn admit(&mut self, entries: Vec, response: AppendResponse) { + // A region would run past the position range of the open batch: the + // batch is sealed and the entries open the next one. The size limit + // seals a batch long before a million entries of one region, so this + // is a theoretical bound. + if !self.open_batch.is_empty() && self.open_batch.would_exhaust_positions(&entries) { + self.flush_open_batch(); + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + return; + } + } + let Some(object_seq) = self.next_object_seq else { + let error = self.poison( + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + .build(), + ); + let _ = response.send(Err(shared(&error))); + return; + }; + let last_entry_ids = match self.open_batch.admit(object_seq, entries) { + Ok(last_entry_ids) => last_entry_ids, + // Nothing was admitted: the append alone runs past the position + // range, which no object can hold. + Err(error) => { + let _ = response.send(Err(error)); + return; + } + }; + for (region_id, entry_id) in &last_entry_ids { + self.issued_entry_ids + .entry(*region_id) + .and_modify(|issued| *issued = (*issued).max(*entry_id)) + .or_insert(*entry_id); + } + match self.ack_mode { + AckMode::Durable => self.pending.push(PendingAppend { + last_entry_ids, + response, + }), + AckMode::Enqueued => { + let _ = response.send(Ok(AppendBatchResponse { last_entry_ids })); + } + } + #[cfg(any(test, feature = "testing"))] + self.admitted_appends.send_modify(|count| *count += 1); + if self.open_batch.should_seal() { + self.flush_open_batch(); + } + } + + /// Returns true once the unpersisted backlog, the open batch and every + /// sealed batch that is not durable, reaches the size or the age threshold. + fn backlog_at_threshold(&self) -> bool { + let bytes = self.open_batch.estimated_bytes() + + self + .sealed + .iter() + .map(|batch| batch.bytes.len()) + .sum::(); + if bytes >= self.max_unpersisted_bytes { + return true; + } + let oldest = self + .sealed + .front() + .map(|batch| batch.first_admitted_at) + .or_else(|| self.open_batch.first_admitted_at()); + oldest.is_some_and(|admitted_at| admitted_at.elapsed() >= self.max_unpersisted_age) + } + + /// Seals the open batch when no create is in flight, so that a stalled + /// append has an upload to wait for. + fn ensure_create_in_flight(&mut self) { + if !self.sealed.iter().any(SealedBatch::is_in_flight) { + self.flush_open_batch(); + } + } + + /// Admits the stalled append once the backlog is below the thresholds and + /// the sealed batches leave room for the two an admission can seal. + fn release_stalled(&mut self) { + if self.stalled.is_none() { + return; + } + if self.is_stopped() { + if let Some((_, response)) = self.take_stalled() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + } + return; + } + if self.backlog_at_threshold() { + // The append needs an upload to complete. + self.ensure_create_in_flight(); + return; + } + if self.sealed.len() + 2 > MAX_SEALED_BATCHES { + return; + } + if let Some((entries, response)) = self.take_stalled() { + self.admit(entries, response); + } + } + + /// Takes the held-back append out on whichever path releases it and + /// records how long it was held back. + fn take_stalled(&mut self) -> Option { + let (append, stalled_at) = self.stalled.take()?; + METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.observe(stalled_at.elapsed().as_secs_f64()); + Some(append) + } + + /// Seals the open batch as the object `next_object_seq` and starts its + /// create when a slot is free. Returns whether a batch was sealed. + fn flush_open_batch(&mut self) -> bool { + if self.open_batch.is_empty() { return false; } - for response in self.stop.drain(..) { + if let Some(error) = terminal(&self.terminal_error) { + self.fail_unacknowledged(|| shared(&error)); + return false; + } + let Some(object_seq) = self.next_object_seq else { + self.poison( + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + .build(), + ); + return false; + }; + + let (entries, first_admitted_at) = self.open_batch.seal(); + let entry_count = entries.len(); + let header = Header { + object_seq, + epoch: self.epoch, + prev: Some( + self.sealed + .back() + .map_or(self.last_indexed, |batch| ChainLink { + object_seq: batch.object_seq, + epoch: self.epoch, + }), + ), + reclaim_boundary: self.reclaim_boundary, + }; + let encoded = match encode_batch(header, entries) { + Ok(encoded) => encoded, + Err(error) => { + self.poison(error); + return false; + } + }; + METRIC_OBJECT_STORE_WAL_OBJECT_BYTES.observe(encoded.bytes.len() as f64); + METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES.observe(entry_count as f64); + self.next_object_seq = object_seq + .checked_add(1) + .filter(|next_object_seq| *next_object_seq < OBJECT_SEQ_LIMIT); + if self.next_object_seq.is_none() { + // The batch takes the last representable sequence: it is created + // and acknowledged, but no later batch can be allocated one. + set_terminal( + &self.terminal_error, + WalObjectSequenceExhaustedSnafu { + last_object_seq: object_seq, + } + .build(), + ); + } + self.sealed.push_back(SealedBatch { + object_seq, + reclaim_boundary: self.reclaim_boundary, + bytes: encoded.bytes, + footer: encoded.footer, + first_admitted_at, + sealed_at: Instant::now(), + attempts: 0, + waiters: std::mem::take(&mut self.pending), + #[cfg(any(test, feature = "testing"))] + seal_waiters: Vec::new(), + state: CreateState::Pending, + }); + self.start_creates(); + true + } + + /// Starts the creates of pending batches in sequence order while fewer + /// than [`MAX_IN_FLIGHT_CREATES`] are in flight, counting the creates of + /// batches that already failed. Nothing starts once stop began, except + /// the backlog of the `enqueued` mode. + fn start_creates(&mut self) { + if self.is_stopped() && self.ack_mode == AckMode::Durable { + return; + } + let mut in_flight = self.creates.len(); + for batch in self.sealed.iter_mut() { + if in_flight >= MAX_IN_FLIGHT_CREATES { + break; + } + if !matches!(batch.state, CreateState::Pending) { + continue; + } + batch.state = CreateState::InFlight; + in_flight += 1; + let delay = if batch.attempts == 0 { + Duration::ZERO + } else { + CREATE_RETRY_DELAY + }; + batch.attempts += 1; + let io = self.io.clone(); + let object_seq = batch.object_seq; + let sealed_at = batch.sealed_at; + self.creating.insert(object_seq); + let bytes = batch.bytes.clone(); + let epoch = self.epoch; + let ack_mode = self.ack_mode; + #[cfg(any(test, feature = "testing"))] + let mut creates_held = self.creates_held.clone(); + #[cfg(any(test, feature = "testing"))] + let creates_fail = self.creates_fail.clone(); + #[cfg(any(test, feature = "testing"))] + let next_create_fails_after_write = self.next_create_fails_after_write.clone(); + #[cfg(any(test, feature = "testing"))] + let parked_creates = self.parked_creates.clone(); + self.creates.push(Box::pin(async move { + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } + // The store was dropped while the create was parked: it never + // runs. + #[cfg(any(test, feature = "testing"))] + if *creates_held.borrow() { + parked_creates.send_modify(|count| *count += 1); + } + #[cfg(any(test, feature = "testing"))] + if creates_held.wait_for(|held| !*held).await.is_err() { + return (object_seq, Err(ObjectStoreWalStoppedSnafu.build())); + } + #[cfg(any(test, feature = "testing"))] + if creates_fail.load(Ordering::Acquire) { + let result = injected_create_failure(io.as_ref(), object_seq); + record_create(&result, Some(sealed_at)); + return (object_seq, result); + } + let created = io.put_if_absent(object_seq, bytes).await; + #[cfg(any(test, feature = "testing"))] + let created = match created { + Ok(_) if next_create_fails_after_write.swap(false, Ordering::AcqRel) => { + injected_create_failure(io.as_ref(), object_seq) + } + created => created, + }; + // Counted as the object store answered, before the epoch of a + // conflicting object is read. + record_create(&created, Some(sealed_at)); + // Any conflict poisons an `enqueued` store, so only the + // `durable` mode reads the epoch of the existing object. + let result = match created { + Err(error @ Error::WalObjectConflict { .. }) + if ack_mode == AckMode::Durable => + { + stale_conflict(io.as_ref(), object_seq, epoch, error).await + } + result => result, + }; + (object_seq, result) + })); + } + } + + fn on_create_completed(&mut self, object_seq: u64, result: Result) { + self.creating.remove(&object_seq); + // A batch that already failed, or that the store gave up on when it + // poisoned itself: the object may exist, but it is off the chain, and + // collection may delete it now. + let Some(index) = self + .sealed + .iter() + .position(|batch| batch.object_seq == object_seq) + else { + self.start_creates(); + self.start_deletes(); + return; + }; + match result { + Ok(_) => self.sealed[index].state = CreateState::Created, + // Nobody is left to retry in the `enqueued` mode, so the store + // repeats a create that failed transiently under the same sequence + // with the same bytes, which an identical retry accepts. + Err(ref error) + if self.ack_mode == AckMode::Enqueued + && !self.is_stopped() + && is_transient(error) => + { + self.sealed[index].state = CreateState::Pending; + } + // The object store did not confirm the object, which may still + // exist or land later, or an earlier epoch holds the sequence and + // can never be on the chain. A caller of the `durable` mode + // retries the append itself; after stop began a transient failure + // drops the `enqueued` backlog. + Err(error @ Error::WalObjectStore { .. }) + if self.ack_mode == AckMode::Durable + || (self.is_stopped() && is_transient(&error)) => + { + self.roll_back(Arc::new(error)) + } + Err(error @ Error::StaleWalObject { .. }) if self.ack_mode == AckMode::Durable => { + self.roll_back(Arc::new(error)) + } + Err(error) => { + self.poison(error); + return; + } + } + self.settle(); + } + + /// Indexes and acknowledges the sealed batches from the front as far as + /// they are created, and starts the creates that a free slot allows. + fn settle(&mut self) { + while let Some(front) = self.sealed.front() { + if !matches!(front.state, CreateState::Created) || !self.index_front() { + break; + } + } + self.start_creates(); + self.release_stalled(); + } + + /// Indexes the created object at the front and acknowledges its waiters. + /// Returns false when the catalog rejected it, which poisons the store. + fn index_front(&mut self) -> bool { + let Some(front) = self.sealed.front() else { + return false; + }; + let indexed = self + .catalog + .write() + .unwrap_or_else(PoisonError::into_inner) + .insert_object(front.object_seq, front.footer.clone()); + if let Err(error) = indexed { + self.poison(error); + return false; + } + let Some(batch) = self.sealed.pop_front() else { + return false; + }; + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.inc(); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.add(batch.bytes.len() as i64); + self.last_indexed = ChainLink { + object_seq: batch.object_seq, + epoch: self.epoch, + }; + // The collection starts before the waiters learn that the object is + // indexed. + self.update_reclaim_boundary(); + if batch.reclaim_boundary > self.durable_reclaim_boundary { + self.durable_reclaim_boundary = batch.reclaim_boundary; + self.collect(); + } + for waiter in batch.waiters { + let _ = waiter.response.send(Ok(AppendBatchResponse { + last_entry_ids: waiter.last_entry_ids, + })); + } + #[cfg(any(test, feature = "testing"))] + for waiter in batch.seal_waiters { + let _ = waiter.send(Ok(())); + } + self.resolve_durable_waiters(); + true + } + + fn resolve_durable_waiters(&mut self) { + let waiters = std::mem::take(&mut self.durable_waiters); + for waiter in waiters { + if self.is_durable_through(waiter.region_id, waiter.entry_id) { + let _ = waiter.response.send(Ok(())); + } else { + self.durable_waiters.push(waiter); + } + } + } + + /// Returns true when no id of the region at or below `entry_id` waits to + /// become durable. An id of a batch that failed will never be durable and + /// no longer waits. + fn is_durable_through(&self, region_id: RegionId, entry_id: EntryId) -> bool { + self.lowest_pending_entry_id(region_id) + .is_none_or(|pending| pending > entry_id) + } + + /// Returns the lowest id of the region that was handed out and is not + /// durable yet. The sealed batches are in sequence order and the open + /// batch is above all of them. + fn lowest_pending_entry_id(&self, region_id: RegionId) -> Option { + self.sealed + .iter() + .flat_map(|batch| batch.footer.iter()) + .find(|entry| entry.region_id == region_id) + .map(|entry| entry.min_entry_id) + .or_else(|| { + self.next_object_seq + .filter(|_| self.open_batch.holds_region(region_id)) + .map(|object_seq| entry_id(object_seq, 1)) + }) + } + + /// Fails every batch that is not indexed after a create failed + /// transiently, including batches already created and batches whose + /// create is in flight. Their sequences are not reused: a create that + /// reported an error may still have written its object. The next batch + /// links to the last indexed object, so no later chain reaches a failed + /// batch. Waiters of a store that was stopped meanwhile learn that + /// instead of the I/O error, like every other entry that never became + /// durable. + fn roll_back(&mut self, error: Arc) { + let stopped = self.is_stopped(); + let failure = || { + if stopped { + ObjectStoreWalStoppedSnafu.build() + } else { + shared(&error) + } + }; + for batch in self.sealed.drain(..) { + self.unindexed.insert(batch.object_seq); + batch.fail(failure); + } + self.reset_open_batch(); + self.fail_unacknowledged(failure); + // An acknowledged backlog was dropped: `stop` reports it, whether + // its caller has arrived yet or not. + if self.ack_mode == AckMode::Enqueued { + self.stop_error.get_or_insert(error); + } + } + + /// Records `error` as terminal and fails every waiter that is not + /// acknowledged with it, or with the stopped error if the store was + /// stopped meanwhile. Creates in flight run to completion, but their + /// outcome is ignored: in the `durable` mode the entries of an object they + /// create were never acknowledged, like those of a crash between creation + /// and acknowledgement; in the `enqueued` mode the acknowledged backlog is + /// discarded and `stop` reports it. Returns the recorded error. + fn poison(&mut self, error: Error) -> Arc { + let error = set_terminal(&self.terminal_error, error); + let stopped = self.is_stopped(); + let failure = || { + if stopped { + ObjectStoreWalStoppedSnafu.build() + } else { + shared(&error) + } + }; + for batch in self.sealed.drain(..) { + batch.fail(failure); + } + self.reset_open_batch(); + self.fail_unacknowledged(failure); + if self.ack_mode == AckMode::Enqueued { + self.stop_error.get_or_insert(error.clone()); + } + error + } + + /// Drops the entries of the open batch; the next admission hands out the + /// same ids again under the sequence the batch is at. + fn reset_open_batch(&mut self) { + self.open_batch.reset(); + } + + /// Fails the waiters of the open batch, the stalled append and the + /// durability waiters. + fn fail_unacknowledged(&mut self, error: impl Fn() -> Error) { + for pending in self.pending.drain(..) { + let _ = pending.response.send(Err(error())); + } + if let Some((_, response)) = self.take_stalled() { + let _ = response.send(Err(error())); + } + for waiter in self.durable_waiters.drain(..) { + let _ = waiter.response.send(Err(error())); + } + } + + fn handle_wait_durable( + &mut self, + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, + ) { + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + return; + } + let durable = { + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + catalog.region_max_entry_id(region_id).unwrap_or(0) + }; + if entry_id <= durable { let _ = response.send(Ok(())); + return; + } + // An acknowledged backlog was dropped: no entry that is not durable + // can be certified any more, whether it was in that backlog or not. + if let Some(error) = &self.stop_error { + let _ = response.send(Err(shared(error))); + return; + } + if self.is_durable_through(region_id, entry_id) { + let _ = response.send(Ok(())); + return; + } + // A caller that stopped waiting leaves a closed response behind. + self.durable_waiters + .retain(|waiter| !waiter.response.is_closed()); + self.durable_waiters.push(DurableWaiter { + region_id, + entry_id, + response, + }); + #[cfg(any(test, feature = "testing"))] + self.durability_waits.send_modify(|count| *count += 1); + } + + /// Makes sure no id of the region at or below `entry_id` is assigned + /// from now on, then records the obsolete watermark of the region; + /// neither is done when the sequence cannot be raised. + fn handle_obsolete( + &mut self, + region_id: RegionId, + entry_id: EntryId, + watermark: EntryId, + response: oneshot::Sender>, + ) { + let result = self.raise_sequence_floor(region_id, entry_id).map(|_| { + record_obsolete(&self.obsolete_entry_ids, region_id, watermark); + }); + if result.is_ok() { + self.update_reclaim_boundary(); + self.collect(); + } + let _ = response.send(result); + } + + /// Raises the reclaim boundary to what the catalog and the obsolete + /// watermarks allow. + fn update_reclaim_boundary(&mut self) { + let boundary = { + let obsolete_entry_ids = self + .obsolete_entry_ids + .lock() + .unwrap_or_else(PoisonError::into_inner); + self.catalog + .read() + .unwrap_or_else(PoisonError::into_inner) + .reclaim_boundary(self.reclaim_boundary, &obsolete_entry_ids) + }; + self.reclaim_boundary = boundary; + } + + /// Starts a collection, which also attempts the deletes that failed since + /// the last one again. + fn collect(&mut self) { + self.failed_deletes.clear(); + self.start_deletes(); + } + + /// Starts deleting the objects below the durable reclaim boundary while + /// fewer than [`MAX_IN_FLIGHT_DELETES`] deletes are in flight, skipping + /// objects whose create is in flight or whose delete failed since the + /// collection started. + fn start_deletes(&mut self) { + if self.is_stopped() || terminal(&self.terminal_error).is_some() { + return; + } + let free = MAX_IN_FLIGHT_DELETES.saturating_sub(self.deletes.len()); + let boundary = self.durable_reclaim_boundary; + let candidates = { + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + catalog + .objects_in_order() + .map(|(object_seq, _)| object_seq) + .take_while(|object_seq| *object_seq < boundary) + .chain(self.unindexed.range(..boundary).copied()) + .filter(|object_seq| { + !self.deleting.contains(*object_seq) + && !self.failed_deletes.contains(object_seq) + && !self.creating.contains(object_seq) + }) + .take(free) + .collect::>() + }; + for object_seq in candidates { + self.deleting.start(object_seq); + let io = self.io.clone(); + self.deletes.push(Box::pin(async move { + (object_seq, io.delete(object_seq).await) + })); + } + } + + /// Unindexes the object `object_seq` and drops its holes once its delete + /// succeeded, then settles its attempt, so a read that the attempt + /// releases sees the outcome in the catalog. + fn on_delete_completed(&mut self, object_seq: u64, result: Result<()>) { + match result { + Ok(()) => { + self.unindexed.remove(&object_seq); + let mut catalog = self.catalog.write().unwrap_or_else(PoisonError::into_inner); + let footer = catalog.remove_object(object_seq); + if let Some(footer) = &footer { + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.dec(); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.sub(object_len(footer) as i64); + } + let mut wal_holes = self + .wal_holes + .lock() + .unwrap_or_else(PoisonError::into_inner); + for entry in footer.iter().flatten() { + if let Some(holes) = wal_holes.get_mut(&entry.region_id) { + holes.remove(&object_seq); + } + } + METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.inc(); + } + Err(error) => { + warn!( + error; + "Failed to delete WAL object {} (sequence {}), it is deleted again at the next collection", + self.io.object_path(object_seq), + object_seq + ); + METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.inc(); + self.failed_deletes.insert(object_seq); + } + } + self.start_deletes(); + self.deleting.finish(object_seq); + } + + /// Moves the sequence of the next object above the object that holds + /// `entry_id` unless it is there already. The move fails while the open + /// batch has handed out ids under the next sequence. A floor that does + /// not fit an entry id poisons the store like an exhausted sequence. + fn raise_sequence_floor(&mut self, region_id: RegionId, entry_id: EntryId) -> Result<()> { + // Nothing is assigned an id after stop began. + if self.is_stopped() { + return Ok(()); + } + if let Some(error) = terminal(&self.terminal_error) { + return Err(shared(&error)); + } + // In the `enqueued` mode an id the store handed out is never handed + // out again while it runs: a create that fails transiently is + // repeated under its sequence and a permanent failure poisons the + // store. Every later id of the region is greater, so the id needs no + // floor even before it is durable. + if self.ack_mode == AckMode::Enqueued + && entry_id <= self.issued_entry_ids.get(®ion_id).copied().unwrap_or(0) + { + return Ok(()); + } + let sequence_floor = sequence_floor(entry_id); + let Some(next_object_seq) = self.next_object_seq else { + return Ok(()); + }; + if next_object_seq >= sequence_floor { + return Ok(()); + } + ensure!( + self.open_batch.is_empty(), + WalObjectSequenceUnsettledSnafu { + object_seq: next_object_seq, + } + ); + if sequence_floor < OBJECT_SEQ_LIMIT { + self.next_object_seq = Some(sequence_floor); + Ok(()) + } else { + self.next_object_seq = None; + let error = self.poison( + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + .build(), + ); + Err(shared(&error)) + } + } + + /// Begins stopping. Nothing is admitted from now on; the `durable` mode + /// drops the open batch and the batches whose create has not started, + /// the `enqueued` mode seals its backlog so it is uploaded. Stop is + /// answered by [`finish_stop`](Self::finish_stop) once nothing is in flight. + fn handle_stop(&mut self, response: oneshot::Sender>) { + self.stop.push(response); + match self.ack_mode { + AckMode::Durable => { + self.reset_open_batch(); + for pending in self.pending.drain(..) { + let _ = pending + .response + .send(Err(ObjectStoreWalStoppedSnafu.build())); + } + if let Some(index) = self + .sealed + .iter() + .position(|batch| matches!(batch.state, CreateState::Pending)) + { + for batch in self.sealed.drain(index..) { + batch.fail(|| ObjectStoreWalStoppedSnafu.build()); + } + } + } + AckMode::Enqueued => { + if let Some((_, response)) = self.take_stalled() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + } + self.flush_open_batch(); + } + } + } + + /// Answers the callers of `stop` once every sealed batch is settled and + /// every create and delete has completed. Returns true when the actor is + /// done. + fn finish_stop(&mut self) -> bool { + if self.stop.is_empty() + || !self.sealed.is_empty() + || !self.creates.is_empty() + || !self.deletes.is_empty() + { + return false; + } + for waiter in self.durable_waiters.drain(..) { + let _ = waiter + .response + .send(Err(ObjectStoreWalStoppedSnafu.build())); + } + let error = self.stop_error.take(); + for response in self.stop.drain(..) { + let _ = response.send(error.as_ref().map_or(Ok(()), |error| Err(shared(error)))); } true } + #[cfg(any(test, feature = "testing"))] + fn handle_seal(&mut self, response: oneshot::Sender>) { + if self.is_stopped() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + return; + } + if self.flush_open_batch() + && let Some(batch) = self.sealed.back_mut() + { + batch.seal_waiters.push(response); + return; + } + let result = terminal(&self.terminal_error).map_or(Ok(()), |error| Err(shared(&error))); + let _ = response.send(result); + } + + /// Fails every caller that waits for the store; the creates are + /// dropped with the actor. + #[cfg(any(test, feature = "testing"))] + fn handle_crash(&mut self) { + let stopped = || ObjectStoreWalStoppedSnafu.build(); + for batch in self.sealed.drain(..) { + batch.fail(stopped); + } + self.reset_open_batch(); + self.fail_unacknowledged(stopped); + for response in self.stop.drain(..) { + let _ = response.send(Err(stopped())); + } + } + fn is_stopped(&self) -> bool { self.stopped.load(Ordering::Acquire) } } +impl Drop for Actor { + /// Settles the deletes the actor still had in flight, so no read waits + /// for an attempt that will never complete, and releases the append still + /// held back, whose caller is gone or about to be. + fn drop(&mut self) { + self.deleting.abandon_all(); + self.take_stalled(); + #[cfg(test)] + LIVE_ACTORS.fetch_sub(1, Ordering::SeqCst); + } +} + +/// Actors not dropped yet. An actor outlives the test that dropped its store +/// and may still move the process-wide metrics. +#[cfg(test)] +static LIVE_ACTORS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + fn terminal(terminal_error: &TerminalError) -> Option> { terminal_error .lock() @@ -405,61 +2042,468 @@ fn terminal(terminal_error: &TerminalError) -> Option> { /// Records `error` as the terminal error unless one is already recorded, and /// returns the recorded one. fn set_terminal(terminal_error: &TerminalError, error: Error) -> Arc { - terminal_error + let mut terminal_error = terminal_error .lock() - .unwrap_or_else(PoisonError::into_inner) + .unwrap_or_else(PoisonError::into_inner); + if terminal_error.is_none() { + METRIC_OBJECT_STORE_WAL_POISONED_TOTAL.inc(); + } + terminal_error .get_or_insert_with(|| Arc::new(error)) .clone() } +/// Counts the outcome of a conditional create and, for the object of a batch +/// sealed at `sealed_at`, times a success. The create of a batch is counted +/// whether or not the batch is still there to be indexed: an object a +/// poisoned store gave up on is as durable as any other. +fn record_create(result: &Result, sealed_at: Option) { + match result { + Ok(_) => { + if let Some(sealed_at) = sealed_at { + METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS + .observe(sealed_at.elapsed().as_secs_f64()); + } + METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL.inc(); + } + Err(Error::WalObjectStore { .. }) => METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL.inc(), + Err(Error::WalObjectConflict { .. }) => { + METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL.inc() + } + Err(_) => {} + } +} + +/// Sets both indexed gauges to the objects `catalog` holds. +fn publish_indexed(catalog: &ObjectCatalog) { + let (objects, bytes) = catalog + .objects_in_order() + .fold((0, 0), |(objects, bytes), (_, footer)| { + (objects + 1, bytes + object_len(footer)) + }); + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.set(objects); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.set(bytes as i64); +} + +/// Returns true for a storage error that a later attempt may not meet: one +/// the object store reports as temporary, or as persistent, which is how its +/// retry layer reports a temporary error that outlasted its retries. +fn is_transient(error: &Error) -> bool { + matches!(error, Error::WalObjectStore { error, .. } if !error.is_permanent()) +} + /// Wraps an error that several callers receive. fn shared(error: &Arc) -> Error { ObjectStoreWalSnafu.into_error(error.clone()) } -/// Rebuilds the catalog from object footers, +/// Returns whether collection deleted the object `object_seq` after a read +/// listed it. The set of deleting objects is observed before the catalog, the +/// reverse of the order the actor updates them, so an object that no longer +/// looks deleting but still looks indexed was never collected. An undecided +/// delete is waited for. +async fn is_collected( + catalog: &RwLock, + deleting: &DeletingObjects, + object_seq: u64, +) -> bool { + if let Some(mut settled) = deleting.attempt(object_seq) { + let _ = settled.wait_for(|settled| *settled).await; + } + !catalog + .read() + .unwrap_or_else(PoisonError::into_inner) + .contains_object(object_seq) +} + +/// Fetches and decodes the segment `entry` describes in the object +/// `object_seq`. A segment that does not decode is fetched once more, since +/// the download rather than the object may be damaged; when it still does not +/// decode the error is returned as [`Error::InvalidWalObject`]. A failed fetch +/// is returned as is. +async fn fetch_segment( + io: &dyn WalObjectIo, + object_seq: u64, + entry: &FooterEntry, +) -> Result> { + let fetch = || io.get_range(object_seq, entry.segment_offset, entry.segment_len); + if let Ok(records) = decode_segment(&fetch().await?, entry) { + return Ok(records); + } + decode_segment(&fetch().await?, entry).with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(object_seq), + }) +} + +/// Records the segment `entry` describes in the object `object_seq` as a hole +/// of its region, once per object, counts the skip and warns about it. An +/// object collected meanwhile is skipped without a hole: the catalog is held +/// while the hole is recorded, as a collection holds it while it drops the +/// holes of the object it removes. +fn record_skipped_segment( + catalog: &RwLock, + wal_holes: &WalHoles, + io: &dyn WalObjectIo, + object_seq: u64, + entry: &FooterEntry, + error: &Error, +) { + let catalog = catalog.read().unwrap_or_else(PoisonError::into_inner); + if !catalog.contains_object(object_seq) { + debug!( + "Skipped WAL object {} (sequence {}), collected after the read listed it: {error}", + io.object_path(object_seq), + object_seq + ); + return; + } + let hole = WalHole { + path: io.object_path(object_seq), + object_seq, + min_entry_id: entry.min_entry_id, + max_entry_id: entry.max_entry_id, + }; + warn!( + error; + "Skipped a corrupted WAL segment of region {}, object {} (sequence {}), entry ids {}..={}", + entry.region_id, + hole.path, + object_seq, + hole.min_entry_id, + hole.max_entry_id + ); + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.inc(); + wal_holes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .entry(entry.region_id) + .or_default() + .entry(object_seq) + .or_insert(hole); +} + +fn encode_batch(header: Header, entries: Vec) -> Result { + let records = entries + .into_iter() + .map(|entry| Record { + region_id: entry.region_id(), + entry_id: entry.entry_id(), + payload: Bytes::from(entry.into_bytes()), + }) + .collect::>(); + encode_object(header, &records) +} + +/// What recovery rebuilt from the objects under a prefix. +#[derive(Debug)] +struct Recovered { + catalog: ObjectCatalog, + /// Sequence above every present object and every id the catalog holds. + next_object_seq: u64, + durable_entry_ids: HashMap, + /// The object the next object extends, `None` on an empty prefix. + tip: Option, + /// Largest epoch any present object carries, zero on an empty prefix. + max_epoch: u64, + /// The reclaim boundary of the tip, zero on an empty prefix. + reclaim_boundary: u64, + /// Present objects that are not indexed, in sequence order. + unindexed: Vec, + /// Number of objects whose header and footer were read. + recovered_objects: u64, +} + +/// The objects recovery fetched, ordered by sequence, and the listed objects +/// that were gone when their header or footer was fetched, with the error the +/// fetch met. A collection of an earlier instance may complete after the +/// listing. +struct FetchedObjects { + objects: Vec, + vanished: Vec<(u64, Error)>, +} + +/// An object as recovery fetched it: its key, header and footer. +struct FetchedObject { + object: ListedObject, + header: Header, + footer: Vec, +} + +/// Rebuilds the catalog from object headers and footers, /// so recovery costs a few small reads per object however large the objects /// are. Segments are not read; a segment checksum is verified by the read that /// decodes it. Footers are fetched for up to [`RECOVERY_CONCURRENCY`] objects /// at a time and indexed in sequence order, so the catalog checks the entry /// ranges of every object against its predecessors like a sequential replay. -async fn recover(io: &dyn WalObjectIo) -> Result<(ObjectCatalog, u64, HashMap)> { +/// Only the objects on the chain [`select_chain`] picks are indexed. +async fn recover(io: &dyn WalObjectIo) -> Result { let objects = io.list().await?; - let mut catalog = ObjectCatalog::default(); - for (object, footer) in fetch_footers(io, objects, RECOVERY_CONCURRENCY).await? { - catalog - .insert_object(object.object_seq, footer) - .with_context(|_| InvalidWalObjectSnafu { path: object.path })?; + finish_recovery(fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?) +} + +/// Indexes the chain of the fetched objects. Objects off the chain, orphans +/// and objects below the reclaim boundary of the tip, are not indexed, but no +/// later object takes their sequences, nor the sequences of vanished objects. +/// A vanished object at or above the boundary of the tip fails recovery with +/// the error its fetch met: collection never deletes such an object. +fn finish_recovery(fetched: FetchedObjects) -> Result { + let FetchedObjects { objects, vanished } = fetched; + let recovered_objects = objects.len() as u64; + let headers = objects + .iter() + .map(|fetched| (fetched.object.object_seq, &fetched.header)) + .collect::>(); + let chain = select_chain(&headers); + // Every open writes an object that starts a chain or extends a complete + // one, and collection removes only objects below the reclaim boundary of + // a later object, so present objects without any complete chain mean + // objects of the chain are gone. + ensure!( + headers.is_empty() || !chain.is_empty(), + CorruptedWalObjectSnafu { + reason: format!( + "no object among {} present objects completes a chain", + headers.len() + ), + } + ); + let tip = chain.last().map(|object_seq| ChainLink { + object_seq: *object_seq, + epoch: headers[object_seq].epoch, + }); + let max_epoch = headers + .values() + .map(|header| header.epoch) + .max() + .unwrap_or(0); + let last_listed = headers + .keys() + .copied() + .chain(vanished.iter().map(|(object_seq, _)| *object_seq)) + .max(); + let after_listed = match last_listed { + None => 0, + Some(last_object_seq) => last_object_seq + .checked_add(1) + .context(WalObjectSequenceExhaustedSnafu { last_object_seq })?, + }; + let reclaim_boundary = tip.map_or(0, |tip| headers[&tip.object_seq].reclaim_boundary); + if let Some((_, error)) = vanished + .into_iter() + .find(|(object_seq, _)| *object_seq >= reclaim_boundary) + { + return Err(error); } - finish_recovery(catalog) -} + let chain = chain.into_iter().collect::>(); -fn finish_recovery( - catalog: ObjectCatalog, -) -> Result<(ObjectCatalog, u64, HashMap)> { - let next_object_seq = catalog.next_object_seq()?; + let mut catalog = ObjectCatalog::default(); + let mut unindexed = Vec::new(); + for FetchedObject { object, footer, .. } in objects { + if chain.contains(&object.object_seq) { + catalog + .insert_object(object.object_seq, footer) + .with_context(|_| InvalidWalObjectSnafu { path: object.path })?; + } else { + unindexed.push(object.object_seq); + } + } + let next_object_seq = catalog.next_object_seq()?.max(after_listed); + ensure!( + next_object_seq < OBJECT_SEQ_LIMIT, + WalObjectSequenceExhaustedSnafu { + last_object_seq: next_object_seq - 1, + } + ); let durable_entry_ids = durable_entry_ids(&catalog); - Ok((catalog, next_object_seq, durable_entry_ids)) + Ok(Recovered { + catalog, + next_object_seq, + durable_entry_ids, + tip, + max_epoch, + reclaim_boundary, + unindexed, + recovered_objects, + }) } -/// Fetches and verifies the footers of `objects`, up to `concurrency` objects -/// at a time, and returns them ordered by object sequence whatever the order -/// the fetches complete in. The first failure abandons the remaining fetches. +/// Returns, in sequence order, the objects on the chain that ends at the tip: +/// the complete object with the largest epoch, then the largest sequence. +/// +/// An object is complete when every link on its chain holds, judged with the +/// reclaim boundary of that object: the chain starts at an object without a +/// predecessor or reaches a predecessor below the boundary, whether or not it +/// is present, and every other link names a present object that carries the +/// epoch the link records. The chain of the tip ends there, so objects below +/// its boundary are not replayed. Only one instance writes +/// under an epoch, so the epoch tells an object of the linking instance from +/// another object under the same sequence. A create that +/// was reported as failed may still leave its object, but every object +/// written after that failure links past it, and every instance writes under +/// an epoch above every object present when it opened, so neither such an +/// object nor a late object of an earlier instance ends the chosen chain. +fn select_chain(headers: &BTreeMap) -> Vec { + let mut candidates = headers.iter().collect::>(); + candidates.sort_unstable_by_key(|(object_seq, header)| Reverse((header.epoch, **object_seq))); + candidates + .into_iter() + .find_map(|(object_seq, _)| chain_ending_at(headers, *object_seq)) + .unwrap_or_default() +} + +/// Returns, in sequence order, the chain that ends at `tip` if `tip` is +/// complete, see [`select_chain`]. +fn chain_ending_at(headers: &BTreeMap, tip: u64) -> Option> { + let reclaim_boundary = headers[&tip].reclaim_boundary; + let mut chain = vec![tip]; + let mut header = headers[&tip]; + while let Some(link) = header.prev { + if link.object_seq >= header.object_seq { + return None; + } + if link.object_seq < reclaim_boundary { + break; + } + header = headers + .get(&link.object_seq) + .filter(|prev| prev.epoch == link.epoch)?; + chain.push(link.object_seq); + } + chain.reverse(); + Some(chain) +} + +/// Writes the empty object that starts the epoch of this instance at +/// `object_seq`, linked to the recovered `tip` and carrying the recovered +/// `reclaim_boundary`, and returns it as the tip later objects extend, +/// together with the sequences of earlier-epoch objects it moved past. The +/// epoch is one above the sequence the create claims, so no two instances +/// share one however their opens interleave, and it is above the epoch of +/// every object recovery listed. +/// +/// An object an earlier instance left at that sequence after recovery listed +/// the prefix has a lower epoch and never ends a chain, so the start object +/// moves to the next sequence and epoch. An object of an equal or later epoch +/// belongs to another writer of the prefix, and the conflict fails the open, +/// as does a create whose outcome is unknown: the next open counts the object +/// in either case. So does a create that finds the same bytes present: +/// another open that recovered the same objects writes an identical start +/// object, so this open cannot claim the epoch even if the object is its own. +async fn start_epoch( + io: &dyn WalObjectIo, + mut object_seq: u64, + tip: Option, + reclaim_boundary: u64, +) -> Result<(ChainLink, Vec)> { + let mut moved_past = Vec::new(); + loop { + ensure!( + object_seq < OBJECT_SEQ_LIMIT, + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + ); + let epoch = object_seq + 1; + let header = Header { + object_seq, + epoch, + prev: tip, + reclaim_boundary, + }; + let created = io + .put_if_absent(object_seq, encode_object(header, &[])?.bytes) + .await; + record_create(&created, None); + match created { + Ok(PutResult::Created) => return Ok((ChainLink { object_seq, epoch }, moved_past)), + Ok(PutResult::AlreadyPresent) => { + return UnconfirmedWalEpochStartSnafu { + path: io.object_path(object_seq), + epoch, + } + .fail(); + } + Err(error @ Error::WalObjectConflict { .. }) => { + if epoch_of(io, object_seq).await? >= epoch { + return Err(error); + } + moved_past.push(object_seq); + object_seq += 1; + } + Err(error) => return Err(error), + } + } +} + +async fn epoch_of(io: &dyn WalObjectIo, object_seq: u64) -> Result { + let head = io.get_range(object_seq, 0, HEADER_LEN as u64).await?; + let header = decode_header(&head).with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(object_seq), + })?; + Ok(header.epoch) +} + +/// Resolves a create of this store's `epoch` that met a different object: +/// an object of an earlier epoch can never be on the chain, so the create +/// fails like a transient error; an object of the same or a later epoch +/// belongs to another writer and `conflict` is returned. +async fn stale_conflict( + io: &dyn WalObjectIo, + object_seq: u64, + epoch: u64, + conflict: Error, +) -> Result { + let existing_epoch = epoch_of(io, object_seq).await?; + if existing_epoch >= epoch { + return Err(conflict); + } + StaleWalObjectSnafu { + path: io.object_path(object_seq), + existing_epoch, + epoch, + } + .fail() +} + +/// Fetches and verifies the headers and footers of `objects`, up to +/// `concurrency` objects at a time, and returns them ordered by object sequence +/// whatever the order the fetches complete in. An object the object store no +/// longer has is returned as vanished; any other failure abandons the +/// remaining fetches. async fn fetch_footers( io: &dyn WalObjectIo, objects: Vec, concurrency: usize, -) -> Result)>> { - let mut footers = futures::stream::iter(objects) +) -> Result { + let outcomes = futures::stream::iter(objects) .map(|object| async move { - let footer = fetch_footer(io, &object).await?; - Ok((object, footer)) + match fetch_footer(io, &object).await { + Ok((header, footer)) => Ok(Ok(FetchedObject { + object, + header, + footer, + })), + Err(error) if is_not_found(&error) => Ok(Err((object.object_seq, error))), + Err(error) => Err(error), + } }) .buffer_unordered(concurrency) .try_collect::>() .await?; - footers.sort_unstable_by_key(|(object, _)| object.object_seq); - Ok(footers) + let mut objects = Vec::with_capacity(outcomes.len()); + let mut vanished = Vec::new(); + for outcome in outcomes { + match outcome { + Ok(object) => objects.push(object), + Err(gone) => vanished.push(gone), + } + } + objects.sort_unstable_by_key(|fetched| fetched.object.object_seq); + Ok(FetchedObjects { objects, vanished }) +} + +fn is_not_found(error: &Error) -> bool { + matches!(error, Error::WalObjectStore { error, .. } if error.kind() == object_store::ErrorKind::NotFound) } /// Reads the header, trailer and footer of `object` and verifies them: the @@ -470,7 +2514,10 @@ async fn fetch_footers( /// A short object is read whole. Otherwise the header and a window at the end /// of the object are read concurrently, and the footer is read separately only /// when it starts before the window. -async fn fetch_footer(io: &dyn WalObjectIo, object: &ListedObject) -> Result> { +async fn fetch_footer( + io: &dyn WalObjectIo, + object: &ListedObject, +) -> Result<(Header, Vec)> { let ListedObject { object_seq, size, .. } = *object; @@ -514,7 +2561,7 @@ async fn fetch_footer(io: &dyn WalObjectIo, object: &ListedObject) -> Result= tail_start { tail.slice(footer_range.start - tail_start..footer_range.end - tail_start) @@ -528,18 +2575,18 @@ async fn fetch_footer(io: &dyn WalObjectIo, object: &ListedObject) -> Result Result<(FixedTrailer, Range)> { +) -> Result<(Header, FixedTrailer, Range)> { let header = decode_header(head)?; ensure!( header.object_seq == object_seq, @@ -552,7 +2599,7 @@ fn locate_footer( ); let trailer = decode_trailer(&tail[tail.len() - TRAILER_LEN..])?; let footer_range = footer_range(trailer, object_len)?; - Ok((trailer, footer_range)) + Ok((header, trailer, footer_range)) } fn durable_entry_ids(catalog: &ObjectCatalog) -> HashMap { @@ -577,6 +2624,8 @@ pub(crate) trait WalObjectIo: Send + Sync { async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result; + async fn delete(&self, object_seq: u64) -> Result<()>; + async fn list(&self) -> Result>; fn object_path(&self, object_seq: u64) -> String; @@ -596,6 +2645,10 @@ impl WalObjectIo for ObjectStoreIo { ObjectStoreIo::get_range(self, object_seq, offset, len).await } + async fn delete(&self, object_seq: u64) -> Result<()> { + ObjectStoreIo::delete(self, object_seq).await + } + async fn list(&self) -> Result> { ObjectStoreIo::list(self).await } @@ -611,18 +2664,22 @@ mod tests { use common_base::readable_size::ReadableSize; use common_error::ext::{ErrorExt, RetryHint}; + use futures::FutureExt; use object_store::services::Memory; + use store_api::logstore::entry::{MultiplePartEntry, MultiplePartHeader}; use tokio::time::timeout; use super::*; use crate::error::WalObjectStoreSnafu; - use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, entry_id}; - use crate::object_store_wal::format::{ - FOOTER_ENTRY_LEN, Header, Record, decode_object, encode_object, - }; + use crate::object_store_wal::batch::{POSITION_LIMIT, entry_id}; + use crate::object_store_wal::format::{FOOTER_ENTRY_LEN, decode_object}; const PREFIX: &str = "wal/datanodes/1/epochs/2"; const WAIT: Duration = Duration::from_secs(30); + /// Held by every test of the module: the metrics of the store are + /// process-wide, so tests running in one process would otherwise move the + /// counters and gauges that another test asserts on. + static SERIALIZED: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); fn memory_store() -> ObjectStore { ObjectStore::new(Memory::default()).unwrap() @@ -637,8 +2694,32 @@ mod tests { } } + /// The same batching as `config`, acknowledging appends on admission. + fn enqueued(config: ObjectStoreWalConfig) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + ack_mode: AckMode::Enqueued, + ..config + } + } + + fn on_corruption( + on_corrupted_segment: CorruptedSegmentAction, + config: ObjectStoreWalConfig, + ) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + on_corrupted_segment, + ..config + } + } + + /// Every append reaches the size limit, so it is persisted on its own. fn eager() -> ObjectStoreWalConfig { - config(Duration::from_millis(10), 1) + config(Duration::from_secs(3600), 1) + } + + /// Nothing is persisted until a test seals the open batch. + fn manual() -> ObjectStoreWalConfig { + config(Duration::from_secs(3600), u64::MAX) } async fn open( @@ -666,8 +2747,137 @@ mod tests { fn id(object_seq: u64, position: u64) -> EntryId { entry_id(object_seq, position) } + + fn entry(store: &ObjectStoreLogStore, region_id: RegionId, data: &str) -> Entry { + store + .entry(data.as_bytes().to_vec(), 0, region_id, &provider(region_id)) + .unwrap() + } + + async fn append( + store: &ObjectStoreLogStore, + region_id: RegionId, + data: &str, + ) -> Result { + store + .append_batch(vec![entry(store, region_id, data)]) + .await + } + + fn spawn_append_batch( + store: &Arc, + entries: Vec, + ) -> tokio::task::JoinHandle> { + let store = store.clone(); + tokio::spawn(async move { store.append_batch(entries).await }) + } + + /// Appends `count` single-entry batches, waiting for each to be admitted + /// before the next is sent, so they are admitted in order. + async fn spawn_appends( + store: &Arc, + region_id: RegionId, + count: usize, + ) -> Vec>> { + let admitted = *store.admitted_appends.borrow(); + let mut handles = Vec::with_capacity(count); + for index in 1..=count { + let entries = vec![entry(store, region_id, &format!("a{index}"))]; + handles.push(spawn_append_batch(store, entries)); + store + .wait_for_admitted_appends(admitted + index) + .await + .unwrap(); + } + handles + } + + async fn object_seqs(io: &dyn WalObjectIo) -> Vec { + io.list() + .await + .unwrap() + .into_iter() + .map(|object| object.object_seq) + .collect() + } + + /// Lists the objects once no delete is in flight. + async fn collected_object_seqs(store: &ObjectStoreLogStore) -> Vec { + timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(); + object_seqs(store.io.as_ref()).await + } + + fn unwrap_shared(error: &Error) -> &Error { + match error { + Error::ObjectStoreWal { source, .. } => source, + other => panic!("expected a shared error, actual {other:?}"), + } + } + + fn assert_stopped(error: &Error) { + assert!( + matches!(error, Error::ObjectStoreWalStopped { .. }), + "unexpected error: {error:?}" + ); + } + + /// Waits for the next operation that `ParkedIo` parked. + async fn next_parked_operation( + parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) -> (u64, oneshot::Sender) { + timeout(WAIT, parked.recv()).await.unwrap().unwrap() + } + + /// Waits for `count` parked operations and returns their releases by + /// sequence. + async fn parked_operations( + parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + count: usize, + ) -> HashMap> { + let mut releases = HashMap::new(); + for _ in 0..count { + let (object_seq, release) = next_parked_operation(parked).await; + releases.insert(object_seq, release); + } + releases + } + + async fn open_over( + io: Arc, + config: &ObjectStoreWalConfig, + ) -> Arc { + ObjectStoreLogStore::open(io, config, PREFIX.to_string()) + .await + .unwrap() + } + + async fn open_parking_creates( + object_store: ObjectStore, + config: &ObjectStoreWalConfig, + ) -> ( + Arc, + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + let (io, parked) = ParkedIo::parking_creates(object_store); + let store = open_over(io.clone(), config).await; + // The start object of the open is created without parking. + io.creates_parked.store(true, Ordering::SeqCst); + (store, io, parked) + } + + /// Round-trips a command through the actor, so an assertion that + /// something did not happen runs after the actor handled every command + /// sent before. The open batch must be empty, as under the eager config. + async fn round_trip_actor(store: &ObjectStoreLogStore) { + store.seal_open_batch().await.unwrap(); + } + #[tokio::test] async fn test_store_rejects_invalid_config() { + let _serialized = SERIALIZED.lock().await; for config in [ config(Duration::from_millis(9), 1), config(Duration::from_secs(1), 0), @@ -675,6 +2885,14 @@ mod tests { prefix: "/absolute".to_string(), ..config(Duration::from_secs(1), 1) }, + ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(0), + ..config(Duration::from_secs(1), 1) + }, + ObjectStoreWalConfig { + max_unpersisted_age: Duration::ZERO, + ..config(Duration::from_secs(1), 1) + }, ] { let error = ObjectStoreLogStore::try_new(memory_store(), &config, 1, 2) .await @@ -689,30 +2907,26 @@ mod tests { #[tokio::test] async fn test_store_recovers_only_its_node_and_generation() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let config = ObjectStoreWalConfig::default(); let identities = [(1, 2), (3, 2), (1, 4)]; for (index, (node_id, generation)) in identities.iter().enumerate() { - let encoded = encode_object( - Header { - object_seq: 0, - writer_instance: [0; 16], - }, + let io = ObjectStoreIo::new( + object_store.clone(), + config.node_prefix(*node_id, *generation), + ) + .unwrap(); + put_fixture( + &io, + 0, &[Record { region_id: region(1), entry_id: index as u64 + 1, payload: Bytes::from_static(b"entry"), }], ) - .unwrap(); - ObjectStoreIo::new( - object_store.clone(), - config.node_prefix(*node_id, *generation), - ) - .unwrap() - .put_if_absent(0, encoded.bytes) - .await - .unwrap(); + .await; } // A root-level object must not be part of any node's recovery. let root_io = ObjectStoreIo::new(object_store.clone(), &config.prefix).unwrap(); @@ -749,31 +2963,49 @@ mod tests { } /// Rebuilds the catalog by decoding whole objects as a recovery oracle. - async fn recover_by_decoding( - io: &dyn WalObjectIo, - ) -> Result<(ObjectCatalog, u64, HashMap)> { - let mut catalog = ObjectCatalog::default(); - for ListedObject { - object_seq, path, .. - } in io.list().await? - { - let bytes = io.get(object_seq).await?; - decode_object(&bytes) + async fn recover_by_decoding(io: &dyn WalObjectIo) -> Result { + let mut objects = Vec::new(); + for object in io.list().await? { + let bytes = io.get(object.object_seq).await?; + let decoded = decode_object(&bytes) .and_then(|decoded| { ensure!( - decoded.header.object_seq == object_seq, + decoded.header.object_seq == object.object_seq, CorruptedWalObjectSnafu { reason: format!( - "header sequence {} does not match key sequence {object_seq}", - decoded.header.object_seq + "header sequence {} does not match key sequence {}", + decoded.header.object_seq, object.object_seq ), } ); - catalog.insert_object(object_seq, decoded.footer) + Ok(decoded) }) - .with_context(|_| InvalidWalObjectSnafu { path })?; + .with_context(|_| InvalidWalObjectSnafu { + path: object.path.clone(), + })?; + objects.push(FetchedObject { + object, + header: decoded.header, + footer: decoded.footer, + }); } - finish_recovery(catalog) + finish_recovery(FetchedObjects { + objects, + vanished: Vec::new(), + }) + } + + fn assert_same_recovery(expected: &Recovered, actual: &Recovered) { + assert_eq!( + catalog_contents(&expected.catalog), + catalog_contents(&actual.catalog) + ); + assert_eq!(expected.next_object_seq, actual.next_object_seq); + assert_eq!(expected.durable_entry_ids, actual.durable_entry_ids); + assert_eq!(expected.tip, actual.tip); + assert_eq!(expected.max_epoch, actual.max_epoch); + assert_eq!(expected.reclaim_boundary, actual.reclaim_boundary); + assert_eq!(expected.unindexed, actual.unindexed); } fn catalog_contents(catalog: &ObjectCatalog) -> Vec<(u64, Vec)> { @@ -798,19 +3030,59 @@ mod tests { } async fn put_records(object_store: &ObjectStore, object_seq: u64, records: &[Record]) { - let encoded = encode_object( - Header { - object_seq, - writer_instance: [0; 16], - }, - records, - ) - .unwrap(); - ObjectStoreIo::new(object_store.clone(), PREFIX) - .unwrap() - .put_if_absent(object_seq, encoded.bytes) + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&io, object_seq, records).await; + } + + /// The header of a fixture object at `object_seq`, which extends the + /// present object right below it under the epoch of that object, so the + /// fixtures a test writes in sequence order form one chain. + async fn fixture_header(io: &ObjectStoreIo, object_seq: u64) -> Header { + let below = io + .list() .await - .unwrap(); + .unwrap() + .into_iter() + .map(|object| object.object_seq) + .filter(|seq| *seq < object_seq) + .max(); + let prev = match below { + Some(seq) => Some(decode_header(&io.get(seq).await.unwrap()).unwrap()), + None => None, + }; + Header { + object_seq, + epoch: prev.as_ref().map_or(1, |prev| prev.epoch), + prev: prev.map(|prev| ChainLink { + object_seq: prev.object_seq, + epoch: prev.epoch, + }), + reclaim_boundary: 0, + } + } + + async fn put_fixture(io: &ObjectStoreIo, object_seq: u64, records: &[Record]) { + let header = fixture_header(io, object_seq).await; + put_object_with_header(io, header, records).await; + } + + async fn put_object_with_header(io: &ObjectStoreIo, header: Header, records: &[Record]) { + let object_seq = header.object_seq; + let encoded = encode_object(header, records).unwrap(); + io.put_if_absent(object_seq, encoded.bytes).await.unwrap(); + } + + /// Puts an object of another writer with `epoch` under `object_seq`. A + /// store opened on an empty prefix writes epoch 1. + async fn put_foreign(object_store: &ObjectStore, object_seq: u64, epoch: u64) { + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let header = Header { + object_seq, + epoch, + prev: None, + reclaim_boundary: 0, + }; + put_object_with_header(&io, header, &[]).await; } fn object_path(object_store: &ObjectStore, object_seq: u64) -> String { @@ -836,6 +3108,18 @@ mod tests { (trailer, footer) } + /// Asserts that recovery over `object_store` rejects the object at `path` + /// for `reason`, whatever a read does with a corrupted segment. + async fn assert_recovery_rejects(object_store: &ObjectStore, path: &str, reason: &str) { + for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { + let config = on_corruption(on_corrupted_segment, eager()); + let error = ObjectStoreLogStore::try_new(object_store.clone(), &config, 1, 2) + .await + .unwrap_err(); + assert_invalid_object(&error, path, reason); + } + } + fn assert_invalid_object(error: &Error, path: &str, reason: &str) { match error { Error::InvalidWalObject { @@ -858,22 +3142,18 @@ mod tests { #[tokio::test] async fn test_store_footer_recovery_matches_full_decode() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); populate(&object_store, 40, 5).await; let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); - let (catalog, next_object_seq, durable) = recover(&io).await.unwrap(); - let (expected_catalog, expected_next_object_seq, expected_durable) = - recover_by_decoding(&io).await.unwrap(); + let recovered = recover(&io).await.unwrap(); + let expected = recover_by_decoding(&io).await.unwrap(); - assert_eq!(40, catalog_contents(&catalog).len()); - assert_eq!( - catalog_contents(&expected_catalog), - catalog_contents(&catalog) - ); - assert_eq!(expected_next_object_seq, next_object_seq); - assert_eq!(40, next_object_seq); - assert_eq!(expected_durable, durable); + assert_eq!(40, catalog_contents(&recovered.catalog).len()); + assert_same_recovery(&expected, &recovered); + assert_eq!(40, recovered.next_object_seq); + let durable = recovered.durable_entry_ids; assert_eq!(5, durable.len()); let store = open(object_store, &eager()).await; @@ -890,6 +3170,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_corrupted_trailer_version_and_footer() { + let _serialized = SERIALIZED.lock().await; type Corrupt = fn(&mut Vec); let cases: &[(&str, Corrupt)] = &[ ("truncated object", |bytes| { @@ -922,22 +3203,17 @@ mod tests { let path = object_path(&object_store, 1); corrupt_object(&object_store, &path, *corrupt).await; - let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) - .await - .unwrap_err(); - assert_invalid_object(&error, &path, reason); + assert_recovery_rejects(&object_store, &path, reason).await; } } #[tokio::test] async fn test_store_recovery_rejects_header_sequence_mismatch() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let encoded = encode_object( - Header { - object_seq: 0, - writer_instance: [0; 16], - }, + fixture_header(&io, 0).await, &[Record { region_id: region(1), entry_id: 1, @@ -947,14 +3223,12 @@ mod tests { .unwrap(); io.put_if_absent(5, encoded.bytes).await.unwrap(); - let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) - .await - .unwrap_err(); - assert_invalid_object( - &error, + assert_recovery_rejects( + &object_store, &io.object_path(5), "header sequence 0 does not match key sequence 5", - ); + ) + .await; } /// Overwrites the byte range of footer entry `index` and refreshes the @@ -972,6 +3246,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_segment_ranges_that_do_not_tile_the_object() { + let _serialized = SERIALIZED.lock().await; type Corrupt = fn(&mut Vec, &[FooterEntry]); let cases: [(&str, Corrupt); 5] = [ ("overflows the object", |bytes, footer| { @@ -1007,10 +3282,7 @@ mod tests { }) .await; - let error = ObjectStoreLogStore::try_new(object_store.clone(), &eager(), 1, 2) - .await - .unwrap_err(); - assert_invalid_object(&error, &path, reason); + assert_recovery_rejects(&object_store, &path, reason).await; let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); let error = recover_by_decoding(&io).await.unwrap_err(); assert_invalid_object(&error, &path, reason); @@ -1018,6 +3290,7 @@ mod tests { } #[tokio::test] async fn test_store_recovery_fetches_a_footer_longer_than_the_tail_window() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let regions = (RECOVERY_TAIL_WINDOW / FOOTER_ENTRY_LEN + 100) as u32; let records = (1..=regions) @@ -1036,16 +3309,10 @@ mod tests { let (trailer, _) = footer_of(&object_store.read(&wide.path).await.unwrap().to_vec()); assert!(trailer.footer_len > RECOVERY_TAIL_WINDOW as u64); - let (catalog, next_object_seq, durable) = recover(io.as_ref()).await.unwrap(); - let (expected_catalog, expected_next_object_seq, expected_durable) = - recover_by_decoding(io.as_ref()).await.unwrap(); - assert_eq!( - catalog_contents(&expected_catalog), - catalog_contents(&catalog) - ); - assert_eq!(expected_next_object_seq, next_object_seq); - assert_eq!(expected_durable, durable); - assert_eq!(regions as usize, durable.len()); + let recovered = recover(io.as_ref()).await.unwrap(); + let expected = recover_by_decoding(io.as_ref()).await.unwrap(); + assert_same_recovery(&expected, &recovered); + assert_eq!(regions as usize, recovered.durable_entry_ids.len()); // The wide object took the header, the tail window and the footer; // the narrow one was read whole. @@ -1096,46 +3363,77 @@ mod tests { payload: Bytes::from(format!("e{entry_id}")), }) .collect::>(); - let encoded = encode_object( - Header { - object_seq, - writer_instance: [0; 16], - }, - &records, - ) - .unwrap(); - ObjectStoreIo::new(object_store.clone(), PREFIX) - .unwrap() - .put_if_absent(object_seq, encoded.bytes) - .await - .unwrap(); + put_records(object_store, object_seq, &records).await; } #[tokio::test] async fn test_store_resumes_sequence_and_durable_ids() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); - let (_, next, durable) = recover(&io).await.unwrap(); - assert_eq!(0, next); - assert!(durable.is_empty()); + let recovered = recover(&io).await.unwrap(); + assert_eq!(0, recovered.next_object_seq); + assert!(recovered.durable_entry_ids.is_empty()); put_object(&object_store, 2, region(1), &[1, id(5, 7)]).await; put_object(&object_store, 4, region(2), &[8]).await; - let (_, next, durable) = recover(&io).await.unwrap(); - assert_eq!(6, next); + let recovered = recover(&io).await.unwrap(); + assert_eq!(6, recovered.next_object_seq); assert_eq!( HashMap::from([(region(1), id(5, 7)), (region(2), 8)]), - durable + recovered.durable_entry_ids ); - let store = open(object_store, &eager()).await; + let store = open(object_store.clone(), &eager()).await; assert_eq!(id(5, 7), latest(&store, region(1))); assert_eq!(8, latest(&store, region(2))); assert_eq!(0, latest(&store, region(3))); + + // The sequence resumes above the object the largest id names, where + // the start object of the store takes object 6, so the first new id + // of every region is greater than every old one. + let response = store + .append_batch(vec![ + entry(&store, region(1), "a"), + entry(&store, region(2), "b"), + ]) + .await + .unwrap(); + assert_eq!( + HashMap::from([(region(1), id(7, 1)), (region(2), id(7, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![2, 4, 6, 7], object_seqs(store.io.as_ref()).await); + assert_eq!( + expected_entries( + region(1), + &[ + (1, "e1"), + (id(5, 7), &format!("e{}", id(5, 7))), + (id(7, 1), "a") + ] + ), + read_entries(&store, region(1), 0).await + ); + store.stop().await.unwrap(); + + // The sequence continues after the start object of the restart. + let store = open(object_store, &eager()).await; + assert_eq!(id(7, 1), latest(&store, region(1))); + let response = append(&store, region(2), "b2").await.unwrap(); + assert_eq!( + HashMap::from([(region(2), id(9, 1))]), + response.last_entry_ids + ); + assert_eq!( + expected_entries(region(2), &[(8, "e8"), (id(7, 1), "b"), (id(9, 1), "b2")]), + read_entries(&store, region(2), 0).await + ); store.stop().await.unwrap(); } #[tokio::test] async fn test_store_rejects_exhausted_sequence() { + let _serialized = SERIALIZED.lock().await; for (object_seq, entry_id) in [(u64::MAX, 1), (OBJECT_SEQ_LIMIT - 1, 1), (0, u64::MAX)] { let object_store = memory_store(); put_object(&object_store, object_seq, region(1), &[entry_id]).await; @@ -1151,6 +3449,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_conflicting_entry_ranges() { + let _serialized = SERIALIZED.lock().await; for entries in [[1, 2], [2, 3]] { let object_store = memory_store(); put_object(&object_store, 0, region(1), &[1, 2]).await; @@ -1164,12 +3463,20 @@ mod tests { #[tokio::test] async fn test_store_rejects_foreign_providers() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let region_id = region(1); let other = Provider::object_store_provider(region_id, "other/prefix".to_string()); let raft = Provider::raft_engine_provider(region_id.as_u64()); for foreign in [&other, &raft] { + let entry = Entry::Naive(NaiveEntry { + provider: foreign.clone(), + region_id, + entry_id: 0, + data: b"a1".to_vec(), + }); let errors = [ + store.append_batch(vec![entry]).await.unwrap_err(), store.latest_entry_id(foreign).unwrap_err(), store.read(foreign, 0, None).await.err().unwrap(), store.create_namespace(foreign).await.unwrap_err(), @@ -1196,6 +3503,7 @@ mod tests { #[tokio::test] async fn test_store_stop_is_idempotent() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let (first, second) = tokio::join!(store.stop(), store.stop()); first.unwrap(); @@ -1203,27 +3511,49 @@ mod tests { timeout(WAIT, store.command_tx.closed()).await.unwrap(); store.stop().await.unwrap(); assert!(store.stopped.load(Ordering::Acquire)); + // Even an empty append fails once stop began. + assert_stopped(&store.append_batch(Vec::new()).await.unwrap_err()); + } + + /// Builds a store whose commands the test receives instead of an actor. + fn store_without_actor() -> (ObjectStoreLogStore, mpsc::Receiver) { + let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); + let (append_tx, _) = mpsc::channel(APPEND_BUFFER); + let store = ObjectStoreLogStore { + prefix: PREFIX.to_string(), + ack_mode: AckMode::Durable, + on_corrupted_segment: CorruptedSegmentAction::Skip, + io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), + catalog: Arc::default(), + obsolete_entry_ids: ObsoleteEntryIds::default(), + wal_holes: WalHoles::default(), + deleting: Arc::default(), + terminal_error: Arc::default(), + stopped: Arc::new(AtomicBool::new(false)), + command_tx, + append_tx, + admitted_appends: watch::channel(0).1, + creates_held: watch::channel(false).0, + creates_fail: Arc::default(), + next_create_fails_after_write: Arc::default(), + parked_creates: watch::channel(0).1, + durability_waits: watch::channel(0).1, + }; + (store, command_rx) } #[tokio::test] async fn test_store_stop_with_actor_gone() { - let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); + let _serialized = SERIALIZED.lock().await; + let (store, command_rx) = store_without_actor(); drop(command_rx); - let store = ObjectStoreLogStore { - prefix: PREFIX.to_string(), - io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), - catalog: Arc::default(), - obsolete_entry_ids: ObsoleteEntryIds::default(), - terminal_error: Arc::default(), - stopped: Arc::new(AtomicBool::new(false)), - command_tx, - }; store.stop().await.unwrap(); store.stop().await.unwrap(); } #[tokio::test] async fn test_store_actor_exits_when_the_store_is_dropped() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let io = Arc::downgrade(&store.io); let stopped = Arc::downgrade(&store.stopped); @@ -1240,6 +3570,7 @@ mod tests { #[tokio::test] async fn test_store_retains_first_terminal_error() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let first = set_terminal( &store.terminal_error, @@ -1253,6 +3584,7 @@ mod tests { let region_id = region(1); let provider = provider(region_id); let errors = [ + store.append_batch(Vec::new()).await.unwrap_err(), store.latest_entry_id(&provider).unwrap_err(), store.read(&provider, 0, None).await.err().unwrap(), store.create_namespace(&provider).await.unwrap_err(), @@ -1276,6 +3608,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_bounds_concurrency_and_orders_by_sequence() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); populate(&object_store, 16, 3).await; let (io, mut parked) = ParkedIo::over(object_store); @@ -1287,8 +3620,9 @@ mod tests { fetch_footers(io.as_ref(), objects, RECOVERY_CONCURRENCY) .await .unwrap() + .objects .into_iter() - .map(|(object, _)| object.object_seq) + .map(|fetched| fetched.object.object_seq) .collect::>() }) }; @@ -1329,6 +3663,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_abandons_pending_fetches_on_failure() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); populate(&object_store, 16, 3).await; let (io, mut parked) = ParkedIo::over(object_store.clone()); @@ -1356,12 +3691,14 @@ mod tests { assert!(parked.try_recv().is_err()); let store = open(object_store, &eager()).await; assert_eq!(id(15, 1), latest(&store, region(1))); - assert_eq!(16, store.catalog.read().unwrap().next_object_seq().unwrap()); + // Sequence 16 holds the object that started the epoch of the store. + assert_eq!(17, store.catalog.read().unwrap().next_object_seq().unwrap()); store.stop().await.unwrap(); } #[tokio::test] async fn test_store_recovery_reads_only_header_and_tail_of_large_object() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); put_records( &object_store, @@ -1399,6 +3736,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_short_read() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); put_object(&object_store, 0, region(1), &[1]).await; let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); @@ -1408,21 +3746,6 @@ mod tests { assert_invalid_object(&error, &object.path, "fewer bytes than the listed"); } - #[tokio::test] - async fn test_store_append_is_unsupported() { - let store = open(memory_store(), &eager()).await; - let error = store.append_batch(Vec::new()).await.unwrap_err(); - assert!(matches!( - error, - Error::UnsupportedObjectStoreWalOperation { .. } - )); - assert_eq!( - common_error::status_code::StatusCode::Unsupported, - error.status_code() - ); - store.stop().await.unwrap(); - } - async fn read_entries( store: &ObjectStoreLogStore, region_id: RegionId, @@ -1456,6 +3779,7 @@ mod tests { #[tokio::test] async fn test_store_reads_only_region_segments_in_entry_order() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let mut expected_ranges = HashMap::<_, Vec<_>>::new(); for (seq, first) in [(1, 10), (3, 30), (7, 70)] { @@ -1530,6 +3854,7 @@ mod tests { #[tokio::test] async fn test_store_obsolete_hides_entries_from_read_only() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); put_object(&object_store, 0, region(1), &[10, 11, 12]).await; put_object(&object_store, 1, region(1), &[20]).await; @@ -1566,10 +3891,7 @@ mod tests { store.stop().await.unwrap(); let reopened = open(object_store, &eager()).await; assert_eq!(all, read_entries(&reopened, region(1), 0).await); - reopened - .obsolete(&p, region(1), EntryId::MAX) - .await - .unwrap(); + reopened.obsolete(&p, region(1), 20).await.unwrap(); assert_eq!( Vec::::new(), read_entries(&reopened, region(1), 0).await @@ -1580,6 +3902,7 @@ mod tests { #[tokio::test] async fn test_store_entry_and_watermarks_reject_another_region() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let p = provider(region(2)); let errors = [ @@ -1608,11 +3931,12 @@ mod tests { store.stop().await.unwrap(); } - #[tokio::test] - async fn test_store_corrupted_segment_fails_the_read_that_decodes_it() { - let object_store = memory_store(); + /// Writes object 0 with one entry of region 1 and one of region 2 and + /// object 1 with an entry of region 2, then flips a byte in the segment of + /// region 2 in object 0. Returns the key of object 0 and its footer. + async fn write_corrupted_segment(object_store: &ObjectStore) -> (String, Vec) { put_records( - &object_store, + object_store, 0, &[ Record { @@ -1628,19 +3952,93 @@ mod tests { ], ) .await; - put_object(&object_store, 1, region(2), &[10]).await; - let path = object_path(&object_store, 0); + put_object(object_store, 1, region(2), &[10]).await; + let path = object_path(object_store, 0); let bytes = object_store.read(&path).await.unwrap().to_vec(); let footer = footer_of(&bytes).1; let corrupt = &footer[1]; - corrupt_object(&object_store, &path, |bytes| { + assert_eq!(region(2), corrupt.region_id); + corrupt_object(object_store, &path, |bytes| { bytes[(corrupt.segment_offset + corrupt.segment_len - 1) as usize] ^= 1; }) .await; + (path, footer) + } + + #[tokio::test] + async fn test_store_skips_a_corrupted_segment_and_records_a_hole() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (path, footer) = write_corrupted_segment(&object_store).await; let (io, reads) = RecordingIo::over(object_store); - let store = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) - .await - .unwrap(); + let store = open_over(io, &eager()).await; + assert_eq!(CorruptedSegmentAction::Skip, store.on_corrupted_segment); + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + // Recovery indexes the object; a read of the other region and one of + // the later object of the damaged region are unaffected. + assert_eq!( + expected_entries(region(1), &[(1, "a1")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 10).await + ); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + + // A read across the damaged segment fetches it twice, continues with + // the next object and records the hole. + reads.lock().unwrap().clear(); + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 0).await + ); + let corrupt = (0, footer[1].segment_offset, footer[1].segment_len); + assert_eq!(corrupt, reads.lock().unwrap()[0]); + assert_eq!(corrupt, reads.lock().unwrap()[1]); + let hole = WalHole { + path, + object_seq: 0, + min_entry_id: 1, + max_entry_id: 1, + }; + assert_eq!( + vec![hole.clone()], + store.wal_holes(&provider(region(2))).unwrap() + ); + assert!(store.wal_holes(&provider(region(1))).unwrap().is_empty()); + assert_eq!( + skipped + 1, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + + // Every read across the segment skips it again; the hole is one. + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 0).await + ); + assert_eq!(vec![hole], store.wal_holes(&provider(region(2))).unwrap()); + assert_eq!( + skipped + 2, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_corrupted_segment_fails_the_read_that_decodes_it() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (path, footer) = write_corrupted_segment(&object_store).await; + let (io, reads) = RecordingIo::over(object_store); + let config = on_corruption(CorruptedSegmentAction::Fail, eager()); + let store = open_over(io, &config).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); assert_eq!(1, latest(&store, region(1))); assert_eq!(10, latest(&store, region(2))); reads.lock().unwrap().clear(); @@ -1669,16 +4067,2820 @@ mod tests { &path, &format!("segment of region {} checksum mismatch", region(2)), ); + let corrupt = (0, footer[1].segment_offset, footer[1].segment_len); + assert_eq!(vec![corrupt, corrupt], *reads.lock().unwrap()); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); assert_eq!( - vec![(0, corrupt.segment_offset, corrupt.segment_len)], - *reads.lock().unwrap() + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() ); store.stop().await.unwrap(); } + #[tokio::test] + async fn test_store_reads_a_segment_whose_first_download_was_damaged() { + let _serialized = SERIALIZED.lock().await; + for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[1, 2]).await; + let (io, reads) = RecordingIo::over(object_store); + let config = on_corruption(on_corrupted_segment, eager()); + let store = open_over(io.clone(), &config).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + // The first download of the segment is damaged, the second is + // not: the read sees the entries and records no hole. + reads.lock().unwrap().clear(); + io.damage_next_range_read.store(true, Ordering::Relaxed); + assert_eq!( + expected_entries(region(1), &[(1, "e1"), (2, "e2")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!(2, reads.lock().unwrap().len()); + assert!(store.wal_holes(&provider(region(1))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_failed_segment_fetch_fails_the_read() { + let _serialized = SERIALIZED.lock().await; + for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[1]).await; + let path = object_path(&object_store, 0); + let (io, _) = RecordingIo::over(object_store.clone()); + let config = on_corruption(on_corrupted_segment, eager()); + let store = open_over(io, &config).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + object_store.delete(&path).await.unwrap(); + let error = store + .read(&provider(region(1)), 0, None) + .await + .unwrap() + .try_collect::>() + .await + .unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { path: actual, .. } if *actual == path), + "unexpected error: {error:?}" + ); + assert!(store.wal_holes(&provider(region(1))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_assigns_ids_from_the_object_sequence() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &manual()).await; + let region_one = region(1); + let region_two = region(2); + + // Appends admitted into one batch share its object; the regions take + // their own positions in admission order, and every append learns + // the last id of each of its regions. + let first = spawn_append_batch( + &store, + vec![ + entry(&store, region_one, "a1"), + entry(&store, region_two, "b1"), + entry(&store, region_one, "a2"), + ], + ); + store.wait_for_admitted_appends(1).await.unwrap(); + let second = spawn_append_batch(&store, vec![entry(&store, region_two, "b2")]); + store.wait_for_admitted_appends(2).await.unwrap(); + assert!(!first.is_finished()); + assert_eq!(0, latest(&store, region_one)); + store.seal_open_batch().await.unwrap(); + assert_eq!( + HashMap::from([(region_one, id(1, 2)), (region_two, id(1, 1))]), + first.await.unwrap().unwrap().last_entry_ids + ); + assert_eq!( + HashMap::from([(region_two, id(1, 2))]), + second.await.unwrap().unwrap().last_entry_ids + ); + + // The next object starts every region at position one again. + let third = spawn_append_batch( + &store, + vec![ + entry(&store, region_two, "b3"), + entry(&store, region_one, "a3"), + ], + ); + store.wait_for_admitted_appends(3).await.unwrap(); + store.seal_open_batch().await.unwrap(); + assert_eq!( + HashMap::from([(region_one, id(2, 1)), (region_two, id(2, 1))]), + third.await.unwrap().unwrap().last_entry_ids + ); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!( + expected_entries( + region_one, + &[(id(1, 1), "a1"), (id(1, 2), "a2"), (id(2, 1), "a3")] + ), + read_entries(&store, region_one, 1).await + ); + assert_eq!( + expected_entries(region_two, &[(id(1, 2), "b2"), (id(2, 1), "b3")]), + read_entries(&store, region_two, id(1, 2)).await + ); + assert_eq!(id(2, 1), latest(&store, region_one)); + assert_eq!(id(2, 1), latest(&store, region_two)); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_seals_when_a_region_exhausts_its_positions() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &manual()).await; + let region_id = region(1); + let entries_of = |count: u64| { + (0..count) + .map(|_| entry(&store, region_id, "")) + .collect::>() + }; + + // The open batch holds every position of the region. + let full = spawn_append_batch(&store, entries_of(POSITION_LIMIT - 1)); + store.wait_for_admitted_appends(1).await.unwrap(); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + + // The next entry of the region seals the batch and opens the next + // object; another region would still have fit. + let next = spawn_append_batch(&store, vec![entry(&store, region_id, "next")]); + store.wait_for_admitted_appends(2).await.unwrap(); + let response = timeout(WAIT, full).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, POSITION_LIMIT - 1))]), + response.last_entry_ids + ); + assert_eq!(vec![1], collected_object_seqs(&store).await); + + // An append that alone runs past the range fits no object: it seals + // the open batch like any append the batch cannot take, then it is + // refused without poisoning the store. + let error = store + .append_batch(entries_of(POSITION_LIMIT)) + .await + .unwrap_err(); + assert!( + matches!(error, Error::WalEntryPositionExhausted { .. }), + "unexpected error: {error:?}" + ); + let response = timeout(WAIT, next).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + let after = spawn_append_batch(&store, vec![entry(&store, region_id, "after")]); + store.wait_for_admitted_appends(3).await.unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, after).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(3, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "next"), (id(3, 1), "after")]), + read_entries(&store, region_id, id(2, 1)).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_flushes_at_the_minimum_interval() { + let _serialized = SERIALIZED.lock().await; + // An append is acknowledged once its object is durable, so a completed + // append proves the tick after it sealed the batch. The appends are + // spaced far apart, so every one of them lands in its own tick and + // the ticks in between find an empty batch. + let interval = MIN_FLUSH_INTERVAL; + let store = open(memory_store(), &config(interval, u64::MAX)).await; + let data = ["a1", "a2", "a3"]; + for (index, data) in data.iter().enumerate() { + tokio::time::sleep(interval * 5).await; + timeout(WAIT, append(&store, region(1), data)) + .await + .unwrap() + .unwrap(); + // Object 1 carries the boundary past the start object. + let expected_seqs = (1..=index as u64 + 1).collect::>(); + assert_eq!(expected_seqs, collected_object_seqs(&store).await); + } + + // Ticks with an empty open batch do not create objects. + tokio::time::sleep(interval * 5).await; + let seqs = collected_object_seqs(&store).await; + assert_eq!(vec![1, 2, 3], seqs); + for object_seq in seqs { + let bytes = store.io.get(object_seq).await.unwrap(); + let decoded = decode_object(&bytes).unwrap(); + assert_eq!(1, decoded.records.len(), "object {object_seq} is empty"); + } + assert_eq!( + expected_entries( + region(1), + &[(id(1, 1), "a1"), (id(2, 1), "a2"), (id(3, 1), "a3")] + ), + read_entries(&store, region(1), 1).await + ); + } + + #[tokio::test] + async fn test_store_rejects_entries_of_another_region() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &manual()).await; + let region_id = region(1); + + let entry = Entry::Naive(NaiveEntry { + provider: provider(region_id), + region_id: region(2), + entry_id: 0, + data: Vec::new(), + }); + let error = store.append_batch(vec![entry]).await.unwrap_err(); + assert!( + matches!(error, Error::MismatchedWalRegion { .. }), + "unexpected error: {error:?}" + ); + let incomplete = Entry::MultiplePart(MultiplePartEntry { + provider: provider(region_id), + region_id, + entry_id: 0, + headers: vec![MultiplePartHeader::First], + parts: vec![b"a1".to_vec()], + }); + let error = store.append_batch(vec![incomplete]).await.unwrap_err(); + assert!( + matches!(error, Error::IncompleteWalEntry { region_id: actual, .. } if actual == region_id), + "unexpected error: {error:?}" + ); + assert_eq!( + common_error::status_code::StatusCode::InvalidArguments, + error.status_code() + ); + assert!( + store + .append_batch(Vec::new()) + .await + .unwrap() + .last_entry_ids + .is_empty() + ); + // Nothing was admitted. + store.seal_open_batch().await.unwrap(); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_pipelines_creates_and_acknowledges_in_sequence_order() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 2).await; + + // At most the limit of creates run at a time; the rest wait for a slot. + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + assert_eq!( + (1..=MAX_IN_FLIGHT_CREATES as u64).collect::>(), + releases.keys().copied().collect::>() + ); + round_trip_actor(&store).await; + assert!(parked.try_recv().is_err()); + assert!(appends.iter().all(|append| !append.is_finished())); + + // Objects 3 and 2 become durable before object 1 and free a slot + // each, but nothing is acknowledged ahead of object 1. + for (released, expected_next) in [(3, 5), (2, 6)] { + releases.remove(&released).unwrap().send(true).unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(expected_next, object_seq); + releases.insert(object_seq, release); + } + round_trip_actor(&store).await; + assert_eq!(vec![0, 2, 3], object_seqs(io.as_ref()).await); + assert!(appends.iter().all(|append| !append.is_finished())); + assert_eq!(0, latest(&store, region_id)); + + // Object 1 releases the acknowledgements of objects 1 to 3. + releases.remove(&1).unwrap().send(true).unwrap(); + let mut appends = appends.into_iter(); + for object_seq in 1..4 { + let response = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + assert_eq!(id(3, 1), latest(&store, region_id)); + + // Object 6 before object 5: the append of object 6 waits for it. + releases.remove(&4).unwrap().send(true).unwrap(); + releases.remove(&6).unwrap().send(true).unwrap(); + let fourth = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(4, 1))]), + fourth.last_entry_ids + ); + let appends = appends.collect::>(); + round_trip_actor(&store).await; + assert!(appends.iter().all(|append| !append.is_finished())); + releases.remove(&5).unwrap().send(true).unwrap(); + for (append, object_seq) in appends.into_iter().zip(5..) { + let response = timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + assert_eq!( + MAX_IN_FLIGHT_CREATES, + io.max_in_flight.load(Ordering::SeqCst) + ); + assert_eq!(vec![1, 2, 3, 4, 5, 6], collected_object_seqs(&store).await); + // Every batch extends the batch sealed before it, whichever was + // durable first. + for object_seq in 1..=6 { + let header = decode_header(&io.get(object_seq).await.unwrap()).unwrap(); + assert_eq!( + Some(object_seq - 1), + header.prev.map(|link| link.object_seq) + ); + } + assert_eq!( + expected_entries( + region_id, + &[ + (id(1, 1), "a1"), + (id(2, 1), "a2"), + (id(3, 1), "a3"), + (id(4, 1), "a4"), + (id(5, 1), "a5"), + (id(6, 1), "a6") + ] + ), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_transient_failure_fails_every_batch_that_is_not_indexed() { + let _serialized = SERIALIZED.lock().await; + let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + // Every slot is taken; the last batch waits for one. + let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 1).await; + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + round_trip_actor(&store).await; + assert!(parked.try_recv().is_err()); + + // Object 1 fails while objects 2 to 4 are in flight and object 5 waits + // for a slot: every batch fails at once. + releases.remove(&1).unwrap().send(false).unwrap(); + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + } + assert_eq!(0, latest(&store, region_id)); + + // The retried entries take the sequences after the last sealed one. + // The creates of the failed batches still hold their slots, so only + // one retry starts until they complete; what they store is off the + // chain. + let retries = spawn_appends(&store, region_id, 2).await; + let (object_seq, first) = next_parked_operation(&mut parked).await; + assert_eq!(6, object_seq); + round_trip_actor(&store).await; + assert!(parked.try_recv().is_err()); + for release in releases.into_values() { + release.send(true).unwrap(); + } + let (object_seq, second) = next_parked_operation(&mut parked).await; + assert_eq!(7, object_seq); + first.send(true).unwrap(); + second.send(true).unwrap(); + for (retry, object_seq) in retries.into_iter().zip(6..) { + let response = timeout(WAIT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + assert_eq!(vec![2, 3, 4, 6, 7], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(6, 1), "a1"), (id(7, 1), "a2")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_transient_failure_before_a_created_object_keeps_it_off_the_chain() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (store, io, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + + // Object 2 is created while object 1 failed: both batches fail, and + // the store keeps serving. + releases.remove(&2).unwrap().send(true).unwrap(); + timeout(WAIT, async { + while object_seqs(io.as_ref()).await != [0, 2] { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + releases.remove(&1).unwrap().send(false).unwrap(); + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + } + assert_eq!(0, latest(&store, region_id)); + + // The retry extends the start object, not the failed batches. + let retry = spawn_appends(&store, region_id, 1).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(3, object_seq); + release.send(true).unwrap(); + let response = timeout(WAIT, retry.into_iter().next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(3, 1))]), + response.last_entry_ids + ); + let bytes = io.get(3).await.unwrap(); + let header = decode_header(&bytes).unwrap(); + assert_eq!( + Some(ChainLink { + object_seq: 0, + epoch: 1, + }), + header.prev + ); + store.stop().await.unwrap(); + + // After a restart only the acknowledged entry replays: object 2 is + // off the chain although it exists. + let store = open(object_store, &eager()).await; + assert_eq!(id(3, 1), latest(&store, region_id)); + assert_eq!( + expected_entries(region_id, &[(id(3, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_permanent_failure_before_a_durable_object_poisons() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + let foreign = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + // Another writer of the same epoch took sequence 1. + put_foreign(&object_store, 1, 1).await; + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + + // Object 2 is created; object 1 conflicts with the foreign object. + releases.remove(&2).unwrap().send(true).unwrap(); + timeout(WAIT, async { + while object_seqs(&foreign).await != [0, 1, 2] { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + round_trip_actor(&store).await; + assert!(appends.iter().all(|append| !append.is_finished())); + releases.remove(&1).unwrap().send(true).unwrap(); + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + } + let error = store.latest_entry_id(&provider(region_id)).unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_conflict_with_an_earlier_epoch_fails_the_batch_without_poisoning() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + // A late object of an earlier epoch lands under sequence 1. + put_foreign(&object_store, 1, 0).await; + let append = spawn_appends(&store, region_id, 1).await; + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + let error = timeout(WAIT, append.into_iter().next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!( + matches!( + unwrap_shared(&error), + Error::StaleWalObject { + existing_epoch: 0, + epoch: 1, + .. + } + ), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + + // The store keeps serving: the retry takes the next sequence. + let retry = spawn_appends(&store, region_id, 1).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(2, object_seq); + release.send(true).unwrap(); + let response = timeout(WAIT, retry.into_iter().next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // The stale object is off the chain after a restart. + let store = open(object_store, &eager()).await; + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_obsolete_raises_the_sequence_floor() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &manual()).await; + let region_one = region(1); + let region_two = region(2); + let mut admitted = 0; + let mut spawn_append = |region_id, data: &str| { + admitted += 1; + ( + spawn_append_batch(&store, vec![entry(&store, region_id, data)]), + admitted, + ) + }; + let (first, count) = spawn_append(region_one, "a1"); + store.wait_for_admitted_appends(count).await.unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, first).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_one, id(1, 1))]), + response.last_entry_ids + ); + let (second, count) = spawn_append(region_one, "a2"); + store.wait_for_admitted_appends(count).await.unwrap(); + + // A durable watermark names an object below the next sequence and + // changes nothing: the open batch goes on under sequence 2. + store + .obsolete(&provider(region_one), region_one, id(1, 1)) + .await + .unwrap(); + let (third, count) = spawn_append(region_one, "a3"); + store.wait_for_admitted_appends(count).await.unwrap(); + + // A watermark that names a later object, as one inherited from + // another prefix does, cannot move the sequence while the open batch + // handed out ids under it: neither the sequence nor the watermark + // moves. + let error = store + .obsolete(&provider(region_two), region_two, id(3, 7)) + .await + .unwrap_err(); + assert!( + matches!( + error, + Error::WalObjectSequenceUnsettled { object_seq: 2, .. } + ), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + assert!( + store + .obsolete_entry_ids + .lock() + .unwrap() + .get(®ion_two) + .is_none() + ); + + // Once the batch is durable the sequence moves above that object, so + // the region's next id is greater than the watermark. + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, second).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_one, id(2, 1))]), + response.last_entry_ids + ); + let response = timeout(WAIT, third).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_one, id(2, 2))]), + response.last_entry_ids + ); + // Object 2 carries the boundary past the watermark of region one. + assert_eq!(vec![2], collected_object_seqs(&store).await); + store + .obsolete(&provider(region_two), region_two, id(3, 7)) + .await + .unwrap(); + assert_eq!( + Some(&id(3, 7)), + store.obsolete_entry_ids.lock().unwrap().get(®ion_two) + ); + let (fourth, count) = spawn_append(region_two, "b1"); + store.wait_for_admitted_appends(count).await.unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, fourth).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_two, id(4, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![2, 4], collected_object_seqs(&store).await); + // The watermark hides nothing of this prefix; the region has no + // entry at or below it. + assert_eq!( + expected_entries(region_two, &[(id(4, 1), "b1")]), + read_entries(&store, region_two, 1).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_does_not_reuse_the_sequence_of_a_failed_create() { + let _serialized = SERIALIZED.lock().await; + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &manual()).await; + let region_one = region(1); + let region_two = region(2); + + // Two appends share object 1, which is written but reported as + // failed: it exists and is not indexed. + let first = spawn_append_batch(&store, vec![entry(&store, region_one, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + let second = spawn_append_batch(&store, vec![entry(&store, region_two, "b1")]); + store.wait_for_admitted_appends(2).await.unwrap(); + io.fail_after_next_put.store(true, Ordering::Relaxed); + store.seal_open_batch().await.unwrap_err(); + for append in [first, second] { + timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + } + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!(0, latest(&store, region_one)); + + // Retried apart, the entries land in later objects rather than + // conflicting with object 1. + for (region_id, data, object_seq) in [(region_one, "a1", 2), (region_two, "b1", 3)] { + let retry = spawn_append_batch(&store, vec![entry(&store, region_id, data)]); + store + .wait_for_admitted_appends(object_seq as usize + 1) + .await + .unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + // Object 3 carries the boundary at object 2, so the failed object 1 + // is collected. + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // Only the retries replay. + let store = open_over(io, &manual()).await; + assert_eq!( + expected_entries(region_one, &[(id(2, 1), "a1")]), + read_entries(&store, region_one, 0).await + ); + assert_eq!( + expected_entries(region_two, &[(id(3, 1), "b1")]), + read_entries(&store, region_two, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_floor_applies_while_a_create_is_in_flight() { + let _serialized = SERIALIZED.lock().await; + let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_one = region(1); + let region_two = region(2); + let pending = spawn_append_batch(&store, vec![entry(&store, region_one, "a1")]); + let (_, release) = next_parked_operation(&mut parked).await; + + // Object 1 is in flight and the open batch is empty: the floor moves + // the next sequence, which a failure of object 1 does not move back. + store + .obsolete(&provider(region_two), region_two, id(5, 1)) + .await + .unwrap(); + release.send(false).unwrap(); + timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err(); + + let write = spawn_append_batch(&store, vec![entry(&store, region_two, "b1")]); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + let response = timeout(WAIT, write).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_two, id(6, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![6], collected_object_seqs(&store).await); + } + + #[tokio::test] + async fn test_store_obsolete_queued_behind_stop_records_the_watermark() { + let _serialized = SERIALIZED.lock().await; + let (store, mut command_rx) = store_without_actor(); + let region_id = region(1); + let provider_one = provider(region_id); + + // The actor exits with the command still queued and never answers + // it: nothing is assigned an id any more, so the watermark alone holds. + let (result, ()) = tokio::join!(store.obsolete(&provider_one, region_id, 1), async { + let command = timeout(WAIT, command_rx.recv()).await.unwrap(); + assert!(matches!( + command, + Some(Command::Obsolete { entry_id: 1, .. }) + )); + }); + result.unwrap(); + assert_eq!( + Some(&1), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + + // The same holds for a call that finds the actor gone. + drop(command_rx); + store + .obsolete(&provider(region(2)), region(2), id(1, 1)) + .await + .unwrap(); + assert_eq!( + Some(&id(1, 1)), + store.obsolete_entry_ids.lock().unwrap().get(®ion(2)) + ); + } + + #[tokio::test] + async fn test_store_hooks_hold_and_fail_creates() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &eager()).await; + let region_id = region(1); + store.hold_creates(); + let held = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + round_trip_actor(&store).await; + assert!(!held.is_finished()); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + store.release_creates(); + timeout(WAIT, held).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + + store.fail_creates(); + let error = append(&store, region_id, "a2").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![1], object_seqs(store.io.as_ref()).await); + + store.begin_stop(); + assert_stopped(&append(&store, region_id, "a3").await.unwrap_err()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_hook_fails_the_next_create_after_it_writes() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &eager()).await; + let region_id = region(1); + + // Object 1 is stored, but its create reports a failure, so the + // append fails and the region has no durable entry. + store.fail_next_create_after_write(); + let error = append(&store, region_id, "a1").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(0, latest(&store, region_id)); + + // Only that create failed. + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_hooks_observe_parked_creates_waits_and_crash() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &enqueued(eager())).await; + let region_id = region(1); + + // The append is acknowledged and its create parks; a wait for its + // entry has to wait. + store.hold_creates(); + append(&store, region_id, "a1").await.unwrap(); + timeout(WAIT, store.wait_for_parked_creates(1)) + .await + .unwrap() + .unwrap(); + let wait = { + let store = store.clone(); + tokio::spawn(async move { store.wait_durable(&provider(region_id), id(1, 1)).await }) + }; + timeout(WAIT, store.wait_for_durability_waits(1)) + .await + .unwrap() + .unwrap(); + assert!(!wait.is_finished()); + + // The crash fails the waiter and the store, and the parked create + // never writes its object. + store.crash().await; + assert_stopped(&timeout(WAIT, wait).await.unwrap().unwrap().unwrap_err()); + assert_stopped(&append(&store, region_id, "a2").await.unwrap_err()); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + } + + #[tokio::test] + async fn test_store_create_held_when_the_store_is_dropped_never_runs() { + let _serialized = SERIALIZED.lock().await; + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &eager()).await; + store.hold_creates(); + let append = spawn_append_batch(&store, vec![entry(&store, region(1), "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + append.abort(); + let _ = append.await; + // A command sender outlives the store, so the actor keeps running and + // the parked create completes on the closed hold channel instead of + // being dropped with the actor. + let command_tx = store.command_tx.clone(); + let terminal_error = store.terminal_error.clone(); + drop(store); + timeout(WAIT, async { + while terminal(&terminal_error).is_none() && object_seqs(io.as_ref()).await == [0] { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + let error = terminal(&terminal_error).unwrap(); + assert!( + matches!(error.as_ref(), Error::ObjectStoreWalStopped { .. }), + "unexpected error: {error:?}" + ); + drop(command_tx); + } + + #[tokio::test] + async fn test_store_rollback_and_poison_drop_the_open_batch() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &manual()).await; + let region_id = region(1); + let spawn_seal = || { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) + }; + + // Object 1 is in flight and a2 waits in the open batch under + // sequence 2 when the create fails: both fail. + let sealed = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + let seal = spawn_seal(); + let (_, release) = next_parked_operation(&mut parked).await; + let open = spawn_append_batch(&store, vec![entry(&store, region_id, "a2")]); + store.wait_for_admitted_appends(2).await.unwrap(); + release.send(false).unwrap(); + for append in [sealed, open] { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + } + timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err(); + + // The open batch had not taken its sequence, so the next entry takes + // the first id of sequence 2 again, alone. + let retry = spawn_append_batch(&store, vec![entry(&store, region_id, "b1")]); + store.wait_for_admitted_appends(3).await.unwrap(); + let seal = spawn_seal(); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + timeout(WAIT, seal).await.unwrap().unwrap().unwrap(); + let response = timeout(WAIT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "b1")]), + read_entries(&store, region_id, 0).await + ); + + // Object 3 conflicts while c2 waits in the open batch: both fail. + let sealed = spawn_append_batch(&store, vec![entry(&store, region_id, "c1")]); + store.wait_for_admitted_appends(4).await.unwrap(); + let seal = spawn_seal(); + let (_, release) = next_parked_operation(&mut parked).await; + let open = spawn_append_batch(&store, vec![entry(&store, region_id, "c2")]); + store.wait_for_admitted_appends(5).await.unwrap(); + put_foreign(&object_store, 3, 1).await; + release.send(true).unwrap(); + for append in [sealed, open] { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + } + timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err(); + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_poisons_once_the_last_sequence_is_taken() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let last = OBJECT_SEQ_LIMIT - 1; + // The start object of the store takes the sequence before the last. + put_object(&object_store, last - 2, region(1), &[id(last - 2, 1)]).await; + let store = open(object_store, &eager()).await; + let next = entry(&store, region(1), "next"); + + // The batch at the last sequence is created and acknowledged, and the + // store is poisoned as soon as it takes that sequence. + let response = append(&store, region(1), "last").await.unwrap(); + assert_eq!( + HashMap::from([(region(1), id(last, 1))]), + response.last_entry_ids + ); + for error in [ + store.latest_entry_id(&provider(region(1))).unwrap_err(), + store.append_batch(vec![next]).await.unwrap_err(), + ] { + assert!( + matches!( + unwrap_shared(&error), + Error::WalObjectSequenceExhausted { last_object_seq, .. } if *last_object_seq == last + ), + "unexpected error: {error:?}" + ); + } + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_holds_appends_back_while_the_sealed_batches_are_full() { + let _serialized = SERIALIZED.lock().await; + let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + let admitted = || *store.admitted_appends.borrow(); + + // Every create is parked: once no two more sealed batches fit, + // further appends wait in the channel and are not admitted. + let mut appends = spawn_appends(&store, region_id, MAX_SEALED_BATCHES - 1).await; + for index in 0..3 { + let entries = vec![entry(&store, region_id, &format!("w{index}"))]; + appends.push(spawn_append_batch(&store, entries)); + } + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + round_trip_actor(&store).await; + assert_eq!(MAX_SEALED_BATCHES - 1, admitted()); + + // A durable object frees a place, and one more append is admitted. + releases.remove(&1).unwrap().send(true).unwrap(); + timeout(WAIT, store.wait_for_admitted_appends(MAX_SEALED_BATCHES)) + .await + .unwrap() + .unwrap(); + round_trip_actor(&store).await; + assert_eq!(MAX_SEALED_BATCHES, admitted()); + + // Stop is handled while appends are held back: the batches whose create + // has not started and the appends still waiting learn of the stop. + let stop = begin_spawned_stop(&store).await; + for release in releases.into_values() { + release.send(true).unwrap(); + } + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + let mut acknowledged = 0; + for append in appends { + match timeout(WAIT, append).await.unwrap().unwrap() { + Ok(_) => acknowledged += 1, + Err(error) => assert_stopped(&error), + } + } + assert_eq!(MAX_IN_FLIGHT_CREATES + 1, acknowledged); + } + + /// Starts `stop` and waits until it has set the stopped flag. + async fn begin_spawned_stop( + store: &Arc, + ) -> tokio::task::JoinHandle> { + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + while !store.stopped.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + stop + } + + #[tokio::test] + async fn test_store_stop_during_failed_flush_reports_stopped() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let pending = spawn_append_batch(&store, vec![entry(&store, region(1), "a1")]); + let (_, release) = next_parked_operation(&mut parked).await; + let stop = begin_spawned_stop(&store).await; + assert!(!stop.is_finished()); + + // The create fails after stop began: its waiter learns of the stop. + release.send(false).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_stopped(&timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err()); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_stopped(&append(&store, region(1), "a2").await.unwrap_err()); + } + + #[tokio::test] + async fn test_store_stop_with_several_creates_in_flight() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + // Four creates are in flight, the fifth batch waits for a slot. + let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 1).await; + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; + assert!(parked.try_recv().is_err()); + + let stop = begin_spawned_stop(&store).await; + assert!(!stop.is_finished()); + // Nothing is assigned an id after stop began, so a watermark needs no + // floor even while creates are in flight. + store + .obsolete(&provider(region(2)), region(2), id(9, 1)) + .await + .unwrap(); + assert_eq!( + Some(&id(9, 1)), + store.obsolete_entry_ids.lock().unwrap().get(®ion(2)) + ); + + // The creates in flight run to completion in any order and are + // acknowledged; the batch that never started learns of the stop. + for object_seq in [3, 1, 4, 2] { + releases.remove(&object_seq).unwrap().send(true).unwrap(); + } + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + let mut appends = appends.into_iter(); + for object_seq in 1..=MAX_IN_FLIGHT_CREATES as u64 { + let response = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + let error = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert_stopped(&error); + // No create was started after stop began. + assert!(parked.try_recv().is_err()); + assert_eq!(vec![0, 1, 2, 3, 4], object_seqs(io.as_ref()).await); + assert_eq!(id(4, 1), latest(&store, region_id)); + timeout(WAIT, store.command_tx.closed()).await.unwrap(); + } + + #[tokio::test] + async fn test_store_stop_during_conflicting_flush_reports_stopped_and_poisons() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + let pending = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + let (_, release) = next_parked_operation(&mut parked).await; + let stop = begin_spawned_stop(&store).await; + // Another writer of the same epoch takes the sequence before the + // create runs. + put_foreign(&object_store, 1, 1).await; + + release.send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_stopped(&timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err()); + for error in [ + store.latest_entry_id(&provider(region_id)).unwrap_err(), + store + .read(&provider(region_id), 1, None) + .await + .err() + .unwrap(), + ] { + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + } + } + + fn injected_failure(operation: &'static str, path: String) -> Result { + Err( + object_store::Error::new(object_store::ErrorKind::Unexpected, "injected failure") + .set_temporary(), + ) + .context(WalObjectStoreSnafu { operation, path }) + } + + fn chain_header(object_seq: u64, epoch: u64, prev: Option<(u64, u64)>) -> Header { + Header { + object_seq, + epoch, + prev: prev.map(|(object_seq, epoch)| ChainLink { object_seq, epoch }), + reclaim_boundary: 0, + } + } + + fn chain_of(headers: &[Header]) -> Vec { + select_chain( + &headers + .iter() + .map(|header| (header.object_seq, header)) + .collect(), + ) + } + + #[test] + fn test_store_selects_the_chain_of_the_latest_complete_object() { + assert!(chain_of(&[]).is_empty()); + // A linear chain is replayed whole. + assert_eq!( + vec![0, 1, 2], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 1, Some((0, 1))), + chain_header(2, 1, Some((1, 1))), + ]) + ); + // Object 1 was reported as failed and object 2 links past it. + assert_eq!( + vec![0, 2], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 1, Some((0, 1))), + chain_header(2, 1, Some((0, 1))), + ]) + ); + // Object 3 extends object 2, which never landed, so object 1 is the + // tip although object 3 has a higher sequence. + assert_eq!( + vec![0, 1], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 1, Some((0, 1))), + chain_header(3, 1, Some((2, 1))), + ]) + ); + // Object 1 links to an object of epoch 2 at sequence 0, where an + // object of epoch 1 is. + assert_eq!( + vec![0], + chain_of(&[chain_header(0, 1, None), chain_header(1, 2, Some((0, 2))),]) + ); + // Instance A left object 2 behind after instance B, a later epoch, + // acknowledged object 1: the later epoch ends the chain. + assert_eq!( + vec![0, 1], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 2, Some((0, 1))), + chain_header(2, 1, Some((0, 1))), + ]) + ); + // A missing predecessor breaks the link, whether or not it lies below + // every present object, so an object that lands late below it cannot + // change which links hold. + assert!( + chain_of(&[ + chain_header(3, 1, Some((2, 1))), + chain_header(4, 1, Some((3, 1))), + ]) + .is_empty() + ); + assert!( + chain_of(&[ + chain_header(2, 1, Some((1, 1))), + chain_header(4, 2, Some((3, 1))), + chain_header(5, 2, Some((4, 2))), + ]) + .is_empty() + ); + assert_eq!( + vec![1], + chain_of(&[chain_header(1, 1, None), chain_header(3, 1, Some((2, 1))),]) + ); + // A predecessor at or above the object itself never holds. + assert!(chain_of(&[chain_header(5, 1, Some((5, 1)))]).is_empty()); + + // A link below the reclaim boundary of the object it is judged for + // holds whether or not the predecessor is present; above it a + // missing predecessor still breaks the link. + let bounded = |object_seq, prev, reclaim_boundary| Header { + reclaim_boundary, + ..chain_header(object_seq, 1, prev) + }; + assert_eq!( + vec![3, 4], + chain_of(&[bounded(3, Some((2, 1)), 3), bounded(4, Some((3, 1)), 3)]) + ); + assert!(chain_of(&[bounded(4, Some((3, 1)), 3)]).is_empty()); + // Object 3 is incomplete with its own boundary, but object 4 only + // needs it to be present. + assert_eq!( + vec![3, 4], + chain_of(&[bounded(3, Some((2, 1)), 0), bounded(4, Some((3, 1)), 3)]) + ); + assert_eq!( + vec![4], + chain_of(&[bounded(3, Some((2, 1)), 0), bounded(4, Some((3, 1)), 4)]) + ); + } + + #[tokio::test] + async fn test_store_recovery_indexes_only_the_chain() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let record = |entry_id: EntryId, data: &'static str| Record { + region_id: region(1), + entry_id, + payload: Bytes::from_static(data.as_bytes()), + }; + put_fixture(&io, 0, &[record(1, "a")]).await; + let first = fixture_header(&io, 1).await; + // Object 1 was reported as failed; object 2 reused its row sequences + // and links past it. + put_object_with_header(&io, first.clone(), &[record(id(1, 1), "old")]).await; + let second = Header { + object_seq: 2, + ..first + }; + put_object_with_header(&io, second, &[record(id(2, 1), "new")]).await; + + for restart in 0..2 { + let store = open(object_store.clone(), &eager()).await; + assert_eq!( + expected_entries(region(1), &[(1, "a"), (id(2, 1), "new")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!(id(2, 1), latest(&store, region(1))); + store.stop().await.unwrap(); + // Every open starts the epoch one above the sequence of its start + // object, and the first restart's start object extends the chain + // the second restart replays. + let start = decode_header(&io.get(3 + restart).await.unwrap()).unwrap(); + assert_eq!(4 + restart, start.epoch); + assert_eq!(Some(2 + restart), start.prev.map(|prev| prev.object_seq)); + } + assert_eq!(vec![0, 1, 2, 3, 4], object_seqs(&io).await); + } + + #[tokio::test] + async fn test_store_recovery_never_reuses_the_sequence_of_an_orphan() { + let _serialized = SERIALIZED.lock().await; + let io = ObjectStoreIo::new(memory_store(), PREFIX).unwrap(); + put_fixture(&io, 0, &[]).await; + // Object 2 extends object 1, which never landed. + let orphan = Header { + prev: Some(ChainLink { + object_seq: 1, + epoch: 1, + }), + ..fixture_header(&io, 2).await + }; + put_object_with_header(&io, orphan, &[]).await; + let recovered = recover(&io).await.unwrap(); + assert_eq!(Some(0), recovered.tip.map(|tip| tip.object_seq)); + assert_eq!(3, recovered.next_object_seq); + } + + #[tokio::test] + async fn test_store_recovery_rejects_objects_without_a_complete_chain() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + // Object 1 extends object 0, which is gone. + put_object_with_header(&io, chain_header(1, 1, Some((0, 1))), &[]).await; + let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::CorruptedWalObject { reason, .. } + if reason.contains("no object among 1 present objects completes a chain")), + "unexpected error: {error:?}" + ); + assert_eq!(vec![1], object_seqs(&io).await); + } + + #[tokio::test] + async fn test_store_recovery_keeps_the_later_epoch_over_a_late_object() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let record = |entry_id: EntryId, data: &'static str| Record { + region_id: region(1), + entry_id, + payload: Bytes::from_static(data.as_bytes()), + }; + put_object_with_header(&io, chain_header(0, 1, None), &[record(1, "a")]).await; + // Instance B acknowledged object 1 before the create instance A issued + // for object 2 landed. + put_object_with_header( + &io, + chain_header(1, 2, Some((0, 1))), + &[record(id(1, 1), "b")], + ) + .await; + put_object_with_header( + &io, + chain_header(2, 1, Some((0, 1))), + &[record(id(2, 1), "late")], + ) + .await; + + let store = open(object_store, &eager()).await; + assert_eq!( + expected_entries(region(1), &[(1, "a"), (id(1, 1), "b")]), + read_entries(&store, region(1), 0).await + ); + store.stop().await.unwrap(); + let start = decode_header(&io.get(3).await.unwrap()).unwrap(); + assert_eq!(4, start.epoch); + assert_eq!( + Some(ChainLink { + object_seq: 1, + epoch: 2, + }), + start.prev + ); + } + + #[tokio::test] + async fn test_store_open_starts_the_first_epoch_on_an_empty_prefix() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + store.stop().await.unwrap(); + let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); + let start = decode_header(&io.get(0).await.unwrap()).unwrap(); + assert_eq!(1, start.epoch); + assert_eq!(None, start.prev); + assert_eq!(vec![0], object_seqs(&io).await); + } + + #[tokio::test] + async fn test_store_start_object_moves_past_an_earlier_epoch_only() { + let _serialized = SERIALIZED.lock().await; + for (epoch, moves) in [(1, true), (2, false)] { + let object_store = memory_store(); + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&fixtures, 0, &[]).await; + // An object lands at the sequence of the start object after + // recovery listed the prefix. + let late = encode_object(chain_header(1, epoch, None), &[]) + .unwrap() + .bytes; + let io = Arc::new(RacingIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + late: Mutex::new(Some((1, late))), + }); + let result = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await; + if moves { + let store = result.unwrap(); + let start = decode_header(&fixtures.get(2).await.unwrap()).unwrap(); + assert_eq!(3, start.epoch); + assert_eq!(Some(0), start.prev.map(|prev| prev.object_seq)); + // Object 3 carries the boundary past the late object, which + // recovery never listed. + append(&store, region(1), "a").await.unwrap(); + assert_eq!(vec![3], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } else { + let error = result.unwrap_err(); + assert!( + matches!(&error, Error::WalObjectConflict { path, .. } if path == &fixtures.object_path(1)), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0, 1], object_seqs(&fixtures).await); + } + } + } + + #[tokio::test] + async fn test_store_open_fails_when_its_start_object_is_already_present() { + let _serialized = SERIALIZED.lock().await; + // On an empty prefix and after object 0 of epoch 1, another open that + // recovered the same objects wrote the very start object this open + // writes. + for present in [false, true] { + let object_store = memory_store(); + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let (start_seq, epoch, tip) = if present { + put_fixture(&fixtures, 0, &[]).await; + (1, 2, Some((0, 1))) + } else { + (0, 1, None) + }; + let same = encode_object(chain_header(start_seq, epoch, tip), &[]) + .unwrap() + .bytes; + let io = Arc::new(RacingIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + late: Mutex::new(Some((start_seq, same))), + }); + let error = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::UnconfirmedWalEpochStart { epoch: actual, path, .. } + if *actual == epoch && path == &fixtures.object_path(start_seq)), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + + // The next open counts that object and starts a later epoch. + open(object_store, &eager()).await.stop().await.unwrap(); + let start = decode_header(&fixtures.get(start_seq + 1).await.unwrap()).unwrap(); + assert_eq!(epoch + 1, start.epoch); + assert_eq!( + Some(ChainLink { + object_seq: start_seq, + epoch, + }), + start.prev + ); + } + } + + #[tokio::test] + async fn test_store_opens_racing_across_a_late_object_take_distinct_epochs() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&io, 0, &[]).await; + // Open A recovers object 0 and pauses before its start object. + let recovered = recover(&io).await.unwrap(); + // A late object of epoch 1 lands at sequence 2, across the gap at + // sequence 1, and open B, which lists it, completes. + put_object_with_header(&io, chain_header(2, 1, Some((0, 1))), &[]).await; + open(object_store, &eager()).await.stop().await.unwrap(); + let b = decode_header(&io.get(3).await.unwrap()).unwrap(); + assert_eq!(4, b.epoch); + + // Open A resumes at the sequence it recovered: the epoch follows from + // the sequence it claims, so it differs from the epoch of open B. + let a = start_epoch( + &io, + recovered.next_object_seq, + recovered.tip, + recovered.reclaim_boundary, + ) + .await + .unwrap() + .0; + assert_eq!( + ChainLink { + object_seq: 1, + epoch: 2, + }, + a + ); + } + + #[tokio::test] + async fn test_store_open_rejects_an_epoch_above_the_next_sequence() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + // No instance that starts at or below sequence 0 writes epoch 5. + put_object_with_header(&io, chain_header(0, 5, None), &[]).await; + let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::CorruptedWalObject { reason, .. } + if reason.contains("carries epoch 5, above the next sequence 1")), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0], object_seqs(&io).await); + } + + #[tokio::test] + async fn test_store_open_fails_when_its_start_object_lands_with_an_unknown_outcome() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&fixtures, 0, &[]).await; + // The start object is stored, but its create reports an error. + let io = Arc::new(LostResponseIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + lose_next: AtomicBool::new(true), + }); + let error = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap_err(); + assert!( + matches!(error, Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + // The open does not move to another sequence. + assert_eq!(vec![0, 1], object_seqs(&fixtures).await); + let stored = decode_header(&fixtures.get(1).await.unwrap()).unwrap(); + assert_eq!(2, stored.epoch); + + // The next open counts the stored start object and claims a later + // epoch. + open(object_store, &eager()).await.stop().await.unwrap(); + let start = decode_header(&fixtures.get(2).await.unwrap()).unwrap(); + assert_eq!(3, start.epoch); + assert_eq!( + Some(ChainLink { + object_seq: 1, + epoch: 2, + }), + start.prev + ); + } + + #[tokio::test] + async fn test_store_enqueued_append_returns_before_the_object_exists() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + + let response = timeout(WAIT, append(&store, region_id, "a1")) + .await + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + assert_eq!(0, latest(&store, region_id)); + assert_eq!(0, store.durable_entry_id(&provider(region_id)).unwrap()); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 2))]), + response.last_entry_ids + ); + + store.seal_open_batch().await.unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!(id(1, 2), latest(&store, region_id)); + assert_eq!( + id(1, 2), + store.durable_entry_id(&provider(region_id)).unwrap() + ); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(1, 2), "a2")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_enqueued_durable_id_advances_after_create_and_indexing() { + let _serialized = SERIALIZED.lock().await; + let (store, _, mut parked) = open_parking_creates(memory_store(), &enqueued(eager())).await; + let region_one = region(1); + let region_two = region(2); + + let response = timeout(WAIT, append(&store, region_one, "a1")) + .await + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_one, id(1, 1))]), + response.last_entry_ids + ); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert_eq!(0, store.durable_entry_id(&provider(region_one)).unwrap()); + let wait = |entry_id| { + let store = store.clone(); + tokio::spawn(async move { store.wait_durable(&provider(region_one), entry_id).await }) + }; + // An id the store never handed out waits for the ids of the region + // that were handed out below it. + let waits = [wait(id(1, 1)), wait(id(1, 7))]; + // The other region has nothing to wait for. + timeout(WAIT, store.wait_durable(&provider(region_two), id(1, 1))) + .await + .unwrap() + .unwrap(); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(waits.iter().all(|wait| !wait.is_finished())); + + release.send(true).unwrap(); + for wait in waits { + timeout(WAIT, wait).await.unwrap().unwrap().unwrap(); + } + assert_eq!( + id(1, 1), + store.durable_entry_id(&provider(region_one)).unwrap() + ); + assert_eq!(0, store.durable_entry_id(&provider(region_two)).unwrap()); + assert_eq!(vec![1], collected_object_seqs(&store).await); + } + + #[tokio::test] + async fn test_store_enqueued_wait_ignores_later_pending_entries_of_the_region() { + let _serialized = SERIALIZED.lock().await; + let (store, _io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let region_a = region(1); + let region_b = region(2); + // Object 1 holds region A's id(1, 1) and object 2 region B's id(2, 1). + for (region_id, data) in [(region_a, "a1"), (region_b, "b1")] { + let response = append(&store, region_id, data).await.unwrap(); + let (_, release) = next_parked_operation(&mut parked).await; + release.send(true).unwrap(); + let entry_id = response.last_entry_ids[®ion_id]; + timeout(WAIT, store.wait_durable(&provider(region_id), entry_id)) + .await + .unwrap() + .unwrap(); + } + let response = append(&store, region_a, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_a, id(3, 1))]), + response.last_entry_ids + ); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(3, object_seq); + + // Every entry of region A up to id(2, 1) is durable, so the wait does + // not depend on the parked create of object 3. + timeout(WAIT, store.wait_durable(&provider(region_a), id(2, 1))) + .await + .unwrap() + .unwrap(); + // A wait for id(3, 1) itself is answered only once object 3 is + // indexed; the round trip orders the check after the actor handled it. + let (response_tx, mut response_rx) = oneshot::channel(); + store + .command_tx + .send(Command::WaitDurable { + region_id: region_a, + entry_id: id(3, 1), + response: response_tx, + }) + .await + .unwrap(); + round_trip_actor(&store).await; + assert!(matches!( + response_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + release.send(true).unwrap(); + timeout(WAIT, response_rx).await.unwrap().unwrap().unwrap(); + } + + #[tokio::test] + async fn test_store_enqueued_obsolete_never_passes_the_durable_id() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + + store + .obsolete(&provider(region_id), region_id, id(1, 2)) + .await + .unwrap(); + assert_eq!( + Some(&0), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + store.seal_open_batch().await.unwrap(); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(1, 2), "a2")]), + read_entries(&store, region_id, 1).await + ); + store + .obsolete(&provider(region_id), region_id, id(1, 2)) + .await + .unwrap(); + assert_eq!( + Some(&id(1, 2)), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + } + + /// Appends one entry in the enqueued mode with `config`, whose create + /// parks, then appends a second one that the backlog threshold must + /// stall. Returns the store, the parked creates, the stalled append and + /// the release of the first create. + async fn stall_second_append( + config: ObjectStoreWalConfig, + ) -> ( + Arc, + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + tokio::task::JoinHandle>, + oneshot::Sender, + ) { + let (store, io, mut parked) = open_parking_creates(memory_store(), &config).await; + let region_id = region(1); + let response = timeout(WAIT, append(&store, region_id, "a1")) + .await + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 1))]), + response.last_entry_ids + ); + if config.max_unpersisted_age < Duration::from_secs(1) { + tokio::time::sleep(config.max_unpersisted_age * 2).await; + } + + let stalled = spawn_append_batch(&store, vec![entry(&store, region_id, "a2")]); + // The stall seals the open batch so that an upload is in flight. + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(!stalled.is_finished()); + assert!(parked.try_recv().is_err()); + (store, io, parked, stalled, release) + } + + #[tokio::test] + async fn test_store_enqueued_backlog_bytes_stall_admission_until_an_upload_completes() { + let _serialized = SERIALIZED.lock().await; + let config = ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(1), + ..enqueued(manual()) + }; + let (store, _, mut parked, stalled, release) = stall_second_append(config).await; + let region_id = region(1); + // Two more appends queue behind the stalled one. + let third = spawn_append_batch(&store, vec![entry(&store, region_id, "a3")]); + let fourth = spawn_append_batch(&store, vec![entry(&store, region_id, "a4")]); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(!third.is_finished() && !fourth.is_finished()); + assert!(parked.try_recv().is_err()); + + // The upload releases the second append, whose entry reaches the + // threshold again; the stall seals it so that the next upload can + // release the third, and so on. + release.send(true).unwrap(); + let response = timeout(WAIT, stalled).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!(id(1, 1), latest(&store, region_id)); + for (append, object_seq) in [(third, 2), (fourth, 3)] { + let (parked_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(object_seq, parked_seq); + release.send(true).unwrap(); + let response = timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq + 1, 1))]), + response.last_entry_ids + ); + } + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + assert_eq!( + expected_entries( + region_id, + &[(id(1, 1), "a1"), (id(2, 1), "a2"), (id(3, 1), "a3")] + ), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_enqueued_backlog_age_stalls_admission_until_an_upload_completes() { + let _serialized = SERIALIZED.lock().await; + let config = ObjectStoreWalConfig { + max_unpersisted_age: Duration::from_millis(50), + ..enqueued(manual()) + }; + let (store, _, _parked, stalled, release) = stall_second_append(config).await; + let region_id = region(1); + + release.send(true).unwrap(); + let response = timeout(WAIT, stalled).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_enqueued_stop_fails_a_stalled_append_and_uploads_the_backlog() { + let _serialized = SERIALIZED.lock().await; + let config = ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(1), + ..enqueued(manual()) + }; + let (store, io, mut parked, stalled, release) = stall_second_append(config).await; + + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + assert_stopped(&timeout(WAIT, stalled).await.unwrap().unwrap().unwrap_err()); + release.send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert!(parked.try_recv().is_err()); + } + + #[tokio::test] + async fn test_store_enqueued_repeats_a_create_that_failed_transiently() { + let _serialized = SERIALIZED.lock().await; + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &enqueued(eager())).await; + let region_id = region(1); + + // The create stores the object but reports a failure; the repeat + // writes the same bytes under the same sequence. + io.fail_after_next_put.store(true, Ordering::Relaxed); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 1))]), + response.last_entry_ids + ); + timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + + // An object store with a retry layer reports a temporary error that + // outlasted its retries as persistent; the create is repeated too. + io.fail_next_put_persistently(); + append(&store, region_id, "a2").await.unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(2, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + } + + #[tokio::test] + async fn test_store_enqueued_permanent_failure_poisons() { + let _serialized = SERIALIZED.lock().await; + // A conflict of any epoch poisons the store without reading the + // existing object, since the acknowledged entries cannot move to + // another sequence. So does a permanent storage error. + for foreign_epoch in [Some(0), Some(2), None] { + let object_store = memory_store(); + let (io, reads) = RecordingIo::over(object_store.clone()); + let store = open_over(io.clone(), &enqueued(eager())).await; + let region_id = region(1); + match foreign_epoch { + Some(epoch) => put_foreign(&object_store, 1, epoch).await, + None => io.fail_next_put_permanently(), + } + let assert_poisoned = |error: Error| { + let poisoned = match foreign_epoch { + Some(_) => matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + None => matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + }; + assert!(poisoned, "unexpected error: {error:?}"); + }; + + // The append was acknowledged; the failure surfaces afterwards. + let second = entry(&store, region_id, "a2"); + append(&store, region_id, "a1").await.unwrap(); + let error = timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap_err(); + assert_poisoned(error); + assert_poisoned(store.append_batch(vec![second]).await.unwrap_err()); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + // The acknowledged entry was dropped, which stop reports. + assert_poisoned(store.stop().await.unwrap_err()); + assert!(reads.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn test_store_enqueued_stop_reports_a_backlog_lost_before_the_stop_command() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + let seal = { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) + }; + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + + // Stop began, but the actor has not received the stop command when + // the create fails: the backlog is dropped and the failure is kept + // for the stop that follows. + store.begin_stop(); + release.send(false).unwrap(); + assert_stopped(&timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err()); + // The lost entry cannot be certified as durable before the stop + // command is handled; a durable id still is. + let error = timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + timeout(WAIT, store.wait_durable(&provider(region_id), 0)) + .await + .unwrap() + .unwrap(); + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert!(parked.try_recv().is_err()); + } + + #[tokio::test] + async fn test_store_enqueued_stop_uploads_the_backlog() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + + store.stop().await.unwrap(); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(id(1, 2), latest(&store, region_id)); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(1, 2), "a2")]), + read_entries(&store, region_id, 1).await + ); + + // A backlog that cannot be uploaded is reported by stop. A transient + // failure, temporary or persistent, drops it; a permanent one also + // poisons the store. + for failure in ["temporary", "persistent", "permanent"] { + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &enqueued(manual())).await; + append(&store, region_id, "a1").await.unwrap(); + match failure { + "temporary" => store.fail_creates(), + "persistent" => io.fail_next_put_persistently(), + _ => io.fail_next_put_permanently(), + } + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_eq!( + failure == "permanent", + store.latest_entry_id(&provider(region_id)).is_err() + ); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_enqueued_issued_id_needs_no_floor_before_it_is_durable() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + + // The open batch holds id(1, 1) under the next sequence. The id was + // handed out, so it is accepted without a floor, with the watermark + // capped to what is durable; an id the store never handed out under + // that sequence is refused. + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + assert_eq!( + Some(&0), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + let error = store + .obsolete(&provider(region_id), region_id, id(1, 2)) + .await + .unwrap_err(); + assert!( + matches!( + error, + Error::WalObjectSequenceUnsettled { object_seq: 1, .. } + ), + "unexpected error: {error:?}" + ); + + store.seal_open_batch().await.unwrap(); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 0).await + ); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + assert!(read_entries(&store, region_id, 0).await.is_empty()); + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + } + + #[tokio::test] + async fn test_store_enqueued_inherited_watermark_raises_the_floor() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let store = open(object_store.clone(), &enqueued(manual())).await; + let region_id = region(1); + + // The region has nothing durable here, so the recorded watermark + // stays at zero, but its ids must still start above the watermark. + store + .obsolete(&provider(region_id), region_id, id(5, 1)) + .await + .unwrap(); + assert_eq!( + Some(&0), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + let response = append(&store, region_id, "r1").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(6, 1))]), + response.last_entry_ids + ); + store.seal_open_batch().await.unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(6, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(id(6, 1), latest(&store, region_id)); + store.stop().await.unwrap(); + + // Replay from the watermark sees the entry. + let store = open(object_store, &enqueued(manual())).await; + assert_eq!( + expected_entries(region_id, &[(id(6, 1), "r1")]), + read_entries(&store, region_id, id(5, 1) + 1).await + ); + } + + #[tokio::test] + async fn test_store_durability_wait_pending_at_stop_fails_with_stopped() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &manual()).await; + let region_id = region(1); + let append = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + let wait = { + let store = store.clone(); + tokio::spawn(async move { store.wait_durable(&provider(region_id), id(1, 1)).await }) + }; + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(!wait.is_finished()); + + store.stop().await.unwrap(); + assert_stopped(&timeout(WAIT, append).await.unwrap().unwrap().unwrap_err()); + assert_stopped(&timeout(WAIT, wait).await.unwrap().unwrap().unwrap_err()); + } + + /// The reclaim boundary in the header of the object `object_seq`. + async fn reclaim_boundary_of(io: &dyn WalObjectIo, object_seq: u64) -> u64 { + decode_header(&io.get(object_seq).await.unwrap()) + .unwrap() + .reclaim_boundary + } + + #[tokio::test] + async fn test_store_stamps_the_reclaim_boundary_and_collects_below_it() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &eager()).await; + let io = store.io.clone(); + let (one, two, three) = (region(1), region(2), region(3)); + let deleted = METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get(); + + // Object 1 carries the boundary past the indexed start object, which + // is collected; region one then holds the boundary at object 1. + append(&store, one, "a1").await.unwrap(); + append(&store, two, "b1").await.unwrap(); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!(1, reclaim_boundary_of(io.as_ref(), 1).await); + assert_eq!(1, reclaim_boundary_of(io.as_ref(), 2).await); + + // A watermark moves the boundary; the next indexed object makes it + // durable. + store.obsolete(&provider(one), one, id(1, 1)).await.unwrap(); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + append(&store, one, "a2").await.unwrap(); + assert_eq!(2, reclaim_boundary_of(io.as_ref(), 3).await); + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); + + // Region one dropped and region two flushed through object 2: object + // 4 carries the sequence after the last indexed object. + store.obsolete(&provider(two), two, id(2, 1)).await.unwrap(); + store.obsolete_all(&provider(one), one).await.unwrap(); + append(&store, two, "b2").await.unwrap(); + assert_eq!(4, reclaim_boundary_of(io.as_ref(), 4).await); + assert_eq!(vec![4], collected_object_seqs(&store).await); + + // A region that appears later does not move the boundary back. + append(&store, three, "c1").await.unwrap(); + assert_eq!(4, reclaim_boundary_of(io.as_ref(), 5).await); + assert_eq!(vec![4, 5], collected_object_seqs(&store).await); + // Other stores of the process count their deletes as well. + assert!(METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get() >= deleted + 4); + assert_eq!( + expected_entries(two, &[(id(4, 1), "b2")]), + read_entries(&store, two, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collects_failed_batches_below_a_durable_boundary() { + let _serialized = SERIALIZED.lock().await; + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &eager()).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(2, 1)) + .await + .unwrap(); + + // Object 3 carries boundary 3 but its create is reported as failed, + // and the create of object 4 fails before it writes: neither boundary + // permits a deletion. + io.fail_after_next_put.store(true, Ordering::Relaxed); + append(&store, region_id, "a3").await.unwrap_err(); + assert_eq!(3, reclaim_boundary_of(io.as_ref(), 3).await); + io.fail_next_put_persistently(); + append(&store, region_id, "a4").await.unwrap_err(); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + + // Object 5 makes boundary 3 durable. The failed batches go once a + // durable boundary passes them, whether their object exists or not. + append(&store, region_id, "a5").await.unwrap(); + assert_eq!(vec![3, 5], collected_object_seqs(&store).await); + store + .obsolete(&provider(region_id), region_id, id(5, 1)) + .await + .unwrap(); + append(&store, region_id, "a6").await.unwrap(); + assert_eq!(vec![6], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(6, 1), "a6")]), + read_entries(&store, region_id, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_attempts_a_failed_delete_again_at_the_next_collection() { + let _serialized = SERIALIZED.lock().await; + let (io, mut parked) = ParkedIo::parking_creates(memory_store()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let region_id = region(1); + let failed = METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.get(); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + release.send(false).unwrap(); + + // The object stays indexed, and an indexing that leaves the durable + // boundary where it is starts no collection. + timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(); + assert_eq!( + failed + 1, + METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.get() + ); + assert!(store.catalog.read().unwrap().contains_object(0)); + append(&store, region_id, "a3").await.unwrap(); + assert!(parked.try_recv().is_err()); + assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + + // The next `obsolete` starts a collection, which deletes it again. + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + release.send(true).unwrap(); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + assert!(!store.catalog.read().unwrap().contains_object(0)); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collection_ends_with_stop_and_poison() { + let _serialized = SERIALIZED.lock().await; + for poison in [false, true] { + let object_store = memory_store(); + let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let region_id = region(1); + for index in 1..=5 { + append(&store, region_id, &format!("a{index}")) + .await + .unwrap(); + } + store + .obsolete(&provider(region_id), region_id, id(5, 1)) + .await + .unwrap(); + append(&store, region_id, "a6").await.unwrap(); + // Objects 0 to 5 are below the durable boundary; four are deleted + // at a time. + let releases = parked_operations(&mut parked, MAX_IN_FLIGHT_DELETES).await; + assert!(parked.try_recv().is_err()); + + let mut stop = None; + if poison { + put_foreign(&object_store, 7, 100).await; + append(&store, region_id, "a7").await.unwrap_err(); + } else { + // The first poll queues the stop command; the actor answers + // the command of `obsolete_all` queued after it only once it + // handled the stop. + let mut stopping = Box::pin(store.stop()); + assert!((&mut stopping).now_or_never().is_none()); + store + .obsolete_all(&provider(region_id), region_id) + .await + .unwrap(); + assert!((&mut stopping).now_or_never().is_none()); + stop = Some(stopping); + } + for release in releases.into_values() { + release.send(true).unwrap(); + } + match stop { + Some(stop) => timeout(WAIT, stop).await.unwrap().unwrap(), + None => timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(), + } + + // The deletes in flight completed and no other started. + assert!(parked.try_recv().is_err()); + assert_eq!( + MAX_IN_FLIGHT_DELETES, + io.max_in_flight.load(Ordering::SeqCst) + ); + let mut expected = vec![4, 5, 6]; + if poison { + expected.push(7); + } + assert_eq!(expected, object_seqs(io.as_ref()).await); + } + } + + #[tokio::test] + async fn test_store_drops_the_holes_of_a_collected_object() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + write_corrupted_segment(&object_store).await; + let store = open(object_store, &eager()).await; + read_entries(&store, region(2), 0).await; + assert_eq!(1, store.wal_holes(&provider(region(2))).unwrap().len()); + + store + .obsolete(&provider(region(1)), region(1), 1) + .await + .unwrap(); + store + .obsolete(&provider(region(2)), region(2), 10) + .await + .unwrap(); + append(&store, region(1), "a").await.unwrap(); + assert_eq!(vec![3], collected_object_seqs(&store).await); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_records_no_hole_for_a_collected_object() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (_, footer) = write_corrupted_segment(&object_store).await; + let store = open(object_store, &eager()).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + // A read met the damaged segment of object 0, which a collection + // removed before the read recorded the hole. + store.catalog.write().unwrap().remove_object(0); + let error = CorruptedWalObjectSnafu { reason: "damaged" }.build(); + record_skipped_segment( + &store.catalog, + &store.wal_holes, + store.io.as_ref(), + 0, + &footer[1], + &error, + ); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collects_a_failed_batch_whose_create_completes_late() { + let _serialized = SERIALIZED.lock().await; + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + // Object 1 fails while the create of object 2 is in flight. + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&1).unwrap().send(false).unwrap(); + for append in appends { + timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + } + + // Objects 3 and 4 make boundary 3 durable; object 2 is skipped while + // its create is in flight. + for object_seq in [3, 4] { + let append = spawn_append_batch(&store, vec![entry(&store, region_id, "a")]); + let (parked_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(object_seq, parked_seq); + release.send(true).unwrap(); + timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + } + assert_eq!(vec![3, 4], collected_object_seqs(&store).await); + + // The create completes and its object is deleted without another + // collection. + io.deletes_parked.store(true, Ordering::SeqCst); + releases.remove(&2).unwrap().send(true).unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(2, object_seq); + release.send(true).unwrap(); + assert_eq!(vec![3, 4], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + + async fn recover_listed(io: &ObjectStoreIo, listed: Vec) -> Result { + finish_recovery(fetch_footers(io, listed, RECOVERY_CONCURRENCY).await?) + } + + #[tokio::test] + async fn test_store_recovery_tolerates_objects_collected_after_the_listing() { + let _serialized = SERIALIZED.lock().await; + let io = ObjectStoreIo::new(memory_store(), PREFIX).unwrap(); + let bounded = |object_seq, prev, reclaim_boundary| Header { + reclaim_boundary, + ..chain_header(object_seq, 1, prev) + }; + put_object_with_header(&io, chain_header(0, 1, None), &[]).await; + put_object_with_header(&io, chain_header(1, 1, Some((0, 1))), &[]).await; + put_object_with_header(&io, bounded(2, Some((1, 1)), 2), &[]).await; + put_object_with_header(&io, bounded(3, Some((2, 1)), 2), &[]).await; + let listed = io.list().await.unwrap(); + + // A delete of an earlier instance removes object 1, below the + // boundary of the tip, after the listing. + io.delete(1).await.unwrap(); + let recovered = recover_listed(&io, listed.clone()).await.unwrap(); + assert_eq!(Some(3), recovered.tip.map(|tip| tip.object_seq)); + assert_eq!(2, recovered.reclaim_boundary); + assert_eq!(vec![0], recovered.unindexed); + assert_eq!(4, recovered.next_object_seq); + + // Object 2 is needed above the boundary: its disappearance fails + // recovery. + io.delete(2).await.unwrap(); + let error = recover_listed(&io, listed).await.unwrap_err(); + assert!(is_not_found(&error), "unexpected error: {error:?}"); + } + + #[tokio::test] + async fn test_store_read_skips_an_object_collected_after_it_was_listed() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &eager()).await; + let region_id = region(1); + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + + // The read lists objects 1 and 2 before a collection deletes object 1. + let stream = store.read(&provider(region_id), 0, None).await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + append(&store, region_id, "a3").await.unwrap(); + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); + + let entries = stream + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect::>(); + assert_eq!(expected_entries(region_id, &[(id(2, 1), "a2")]), entries); + assert!(store.wal_holes(&provider(region_id)).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_read_waits_for_an_undecided_delete() { + let _serialized = SERIALIZED.lock().await; + for deleted in [true, false] { + let object_store = memory_store(); + let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let region_id = region(1); + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + append(&store, region_id, "a1").await.unwrap(); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + append(&store, region_id, "a2").await.unwrap(); + + // The read lists objects 1 and 2. The delete of object 1 removes + // it, but its outcome is not reported yet. + let stream = store.read(&provider(region_id), 0, None).await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + append(&store, region_id, "a3").await.unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + object_store + .delete(&object_path(&object_store, 1)) + .await + .unwrap(); + let read = tokio::spawn(stream.try_collect::>()); + let waiting = || { + store + .deleting + .lock() + .get(&1) + .map_or(0, |settled| settled.receiver_count()) + }; + timeout(WAIT, async { + while waiting() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(!read.is_finished()); + + // A delete that succeeded skips the object; one that failed + // leaves it indexed, and the read fails on the fetch. + release.send(deleted).unwrap(); + let result = timeout(WAIT, read).await.unwrap().unwrap(); + if deleted { + let entries = result.unwrap().into_iter().flatten().collect::>(); + assert_eq!(expected_entries(region_id, &[(id(2, 1), "a2")]), entries); + } else { + let error = result.unwrap_err(); + assert!(is_not_found(&error), "unexpected error: {error:?}"); + } + assert!(store.wal_holes(&provider(region_id)).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_recovers_from_the_reclaim_boundary_of_the_tip() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let region_id = region(1); + let store = open(object_store.clone(), &eager()).await; + append(&store, region_id, "a1").await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + append(&store, region_id, "a2").await.unwrap(); + assert_eq!(vec![2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // Object 2 links to object 1, below its boundary. The start object + // carries the recovered boundary. + let store = open(object_store.clone(), &eager()).await; + assert_eq!(2, reclaim_boundary_of(&io, 3).await); + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "a2")]), + read_entries(&store, region_id, 0).await + ); + store + .obsolete(&provider(region_id), region_id, id(2, 1)) + .await + .unwrap(); + append(&store, region_id, "a3").await.unwrap(); + assert_eq!(vec![4], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + let store = open(object_store, &eager()).await; + assert_eq!(4, reclaim_boundary_of(&io, 5).await); + assert_eq!( + expected_entries(region_id, &[(id(4, 1), "a3")]), + read_entries(&store, region_id, 0).await + ); + assert_eq!(vec![4, 5], object_seqs(&io).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collects_a_late_object_of_an_earlier_epoch_below_the_boundary() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let region_id = region(1); + let record = |entry_id, data: &'static str| Record { + region_id, + entry_id, + payload: Bytes::from_static(data.as_bytes()), + }; + let bounded = |object_seq, epoch, prev, reclaim_boundary| Header { + reclaim_boundary, + ..chain_header(object_seq, epoch, prev) + }; + put_object_with_header(&io, chain_header(0, 1, None), &[]).await; + put_object_with_header( + &io, + chain_header(1, 1, Some((0, 1))), + &[record(id(1, 1), "a")], + ) + .await; + // Epoch 4 started at object 3, and object 4 carries boundary 4. + // Object 2 of epoch 1 landed late; its link holds, but its epoch is + // lower. + put_object_with_header( + &io, + bounded(2, 1, Some((1, 1)), 2), + &[record(id(2, 1), "late")], + ) + .await; + put_object_with_header(&io, chain_header(3, 4, Some((1, 1))), &[]).await; + put_object_with_header( + &io, + bounded(4, 4, Some((3, 4)), 4), + &[record(id(4, 1), "b")], + ) + .await; + let recovered = recover(&io).await.unwrap(); + assert_eq!(Some(4), recovered.tip.map(|tip| tip.object_seq)); + assert_eq!(4, recovered.reclaim_boundary); + assert_eq!(vec![0, 1, 2, 3], recovered.unindexed); + + // The start object makes boundary 4 durable, and the first + // collection deletes everything below it. + let store = open(object_store, &eager()).await; + assert_eq!( + expected_entries(region_id, &[(id(4, 1), "b")]), + read_entries(&store, region_id, 0).await + ); + store + .obsolete(&provider(region_id), region_id, id(4, 1)) + .await + .unwrap(); + assert_eq!(vec![4, 5], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + + /// Object access whose next create stores the object but reports a + /// transient failure, as a create whose response is lost. + struct LostResponseIo { + inner: ObjectStoreIo, + lose_next: AtomicBool, + } + + #[async_trait::async_trait] + impl WalObjectIo for LostResponseIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + let result = self.inner.put_if_absent(object_seq, content).await?; + if !self.lose_next.swap(false, Ordering::SeqCst) { + return Ok(result); + } + Err( + object_store::Error::new(object_store::ErrorKind::Unexpected, "lost response") + .set_temporary(), + ) + .context(WalObjectStoreSnafu { + operation: "write", + path: self.object_path(object_seq), + }) + } + + async fn get(&self, object_seq: u64) -> Result { + self.inner.get(object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + self.inner.get_range(object_seq, offset, len).await + } + + async fn delete(&self, object_seq: u64) -> Result<()> { + self.inner.delete(object_seq).await + } + + async fn list(&self) -> Result> { + self.inner.list().await + } + + fn object_path(&self, object_seq: u64) -> String { + self.inner.object_path(object_seq) + } + } + + /// Object access whose first create lets another object land at a given + /// sequence first, as a create issued before recovery would. + struct RacingIo { + inner: ObjectStoreIo, + late: Mutex>, + } + + #[async_trait::async_trait] + impl WalObjectIo for RacingIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + let late = self.late.lock().unwrap().take(); + if let Some((late_seq, late)) = late { + self.inner.put_if_absent(late_seq, late).await?; + } + self.inner.put_if_absent(object_seq, content).await + } + + async fn get(&self, object_seq: u64) -> Result { + self.inner.get(object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + self.inner.get_range(object_seq, offset, len).await + } + + async fn delete(&self, object_seq: u64) -> Result<()> { + self.inner.delete(object_seq).await + } + + async fn list(&self) -> Result> { + self.inner.list().await + } + + fn object_path(&self, object_seq: u64) -> String { + self.inner.object_path(object_seq) + } + } + + /// Object access whose whole-object reads, or conditional creates, and + /// on request deletes, park until the test releases them, counting how + /// many are in flight. A release of false fails the operation with a + /// transient error before it reaches the object store. struct ParkedIo { inner: ObjectStoreIo, parked: mpsc::UnboundedSender<(u64, oneshot::Sender)>, + park_creates: bool, + /// Whether creates park yet; set once the store opened. + creates_parked: AtomicBool, + /// Whether deletes park. + deletes_parked: AtomicBool, in_flight: AtomicUsize, max_in_flight: AtomicUsize, } @@ -1689,43 +6891,69 @@ mod tests { ) -> ( Arc, mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + Self::new(object_store, false) + } + + fn parking_creates( + object_store: ObjectStore, + ) -> ( + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + Self::new(object_store, true) + } + + fn new( + object_store: ObjectStore, + park_creates: bool, + ) -> ( + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, ) { let (parked, parked_rx) = mpsc::unbounded_channel(); ( Arc::new(Self { inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), parked, + park_creates, + creates_parked: AtomicBool::new(false), + deletes_parked: AtomicBool::new(false), in_flight: AtomicUsize::new(0), max_in_flight: AtomicUsize::new(0), }), parked_rx, ) } + + /// Parks until the test releases `object_seq` and returns whether the + /// operation proceeds. + async fn park(&self, object_seq: u64) -> bool { + let in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(in_flight, Ordering::SeqCst); + let (release, released) = oneshot::channel(); + self.parked.send((object_seq, release)).unwrap(); + let proceed = released.await.unwrap(); + self.in_flight.fetch_sub(1, Ordering::SeqCst); + proceed + } } #[async_trait::async_trait] impl WalObjectIo for ParkedIo { async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + if self.park_creates + && self.creates_parked.load(Ordering::SeqCst) + && !self.park(object_seq).await + { + return injected_failure("write", self.object_path(object_seq)); + } self.inner.put_if_absent(object_seq, content).await } async fn get(&self, object_seq: u64) -> Result { - let in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.max_in_flight.fetch_max(in_flight, Ordering::SeqCst); - let (release, released) = oneshot::channel(); - self.parked.send((object_seq, release)).unwrap(); - let success = released.await.unwrap(); - self.in_flight.fetch_sub(1, Ordering::SeqCst); - if !success { - return Err(object_store::Error::new( - object_store::ErrorKind::Unexpected, - "injected failure", - ) - .set_temporary()) - .context(WalObjectStoreSnafu { - operation: "read", - path: self.object_path(object_seq), - }); + if !self.park_creates && !self.park(object_seq).await { + return injected_failure("read", self.object_path(object_seq)); } self.inner.get(object_seq).await } @@ -1734,6 +6962,13 @@ mod tests { self.inner.get_range(object_seq, offset, len).await } + async fn delete(&self, object_seq: u64) -> Result<()> { + if self.deletes_parked.load(Ordering::SeqCst) && !self.park(object_seq).await { + return injected_failure("delete", self.object_path(object_seq)); + } + self.inner.delete(object_seq).await + } + async fn list(&self) -> Result> { self.inner.list().await } @@ -1745,10 +6980,18 @@ mod tests { type RangeReads = Arc>>; - /// Object access that records every range read as (sequence, offset, length). + /// Object access that records every range read as (sequence, offset, + /// length) and, on request, reports the next conditional create as failed + /// after it wrote the object, fails it with a given error before it + /// writes, stores the same content before it writes, or damages the bytes + /// the next range read returns. struct RecordingIo { inner: ObjectStoreIo, reads: RangeReads, + fail_after_next_put: AtomicBool, + fail_next_put: Mutex>, + damage_next_range_read: AtomicBool, + duplicate_next_put: AtomicBool, } impl RecordingIo { @@ -1757,15 +7000,51 @@ mod tests { let io = Self { inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), reads: reads.clone(), + fail_after_next_put: AtomicBool::new(false), + fail_next_put: Mutex::new(None), + damage_next_range_read: AtomicBool::new(false), + duplicate_next_put: AtomicBool::new(false), }; (Arc::new(io), reads) } + + fn fail_next_put_permanently(&self) { + let error = object_store::Error::new( + object_store::ErrorKind::PermissionDenied, + "injected failure", + ); + *self.fail_next_put.lock().unwrap() = Some(error); + } + + /// Fails the next create as a retry layer reports a temporary error + /// that outlasted its retries. + fn fail_next_put_persistently(&self) { + let error = object_store::Error::new(object_store::ErrorKind::Unexpected, "injected") + .set_temporary() + .set_persistent(); + *self.fail_next_put.lock().unwrap() = Some(error); + } } #[async_trait::async_trait] impl WalObjectIo for RecordingIo { async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { - self.inner.put_if_absent(object_seq, content).await + if let Some(error) = self.fail_next_put.lock().unwrap().take() { + return Err(error).context(WalObjectStoreSnafu { + operation: "write", + path: self.object_path(object_seq), + }); + } + if self.duplicate_next_put.swap(false, Ordering::Relaxed) { + self.inner + .put_if_absent(object_seq, content.clone()) + .await?; + } + let result = self.inner.put_if_absent(object_seq, content).await?; + if self.fail_after_next_put.swap(false, Ordering::Relaxed) { + return injected_failure("write", self.object_path(object_seq)); + } + Ok(result) } async fn get(&self, object_seq: u64) -> Result { @@ -1774,7 +7053,19 @@ mod tests { async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { self.reads.lock().unwrap().push((object_seq, offset, len)); - self.inner.get_range(object_seq, offset, len).await + let bytes = self.inner.get_range(object_seq, offset, len).await?; + if self.damage_next_range_read.swap(false, Ordering::Relaxed) { + let mut damaged = bytes.to_vec(); + if let Some(last) = damaged.last_mut() { + *last ^= 1; + } + return Ok(Bytes::from(damaged)); + } + Ok(bytes) + } + + async fn delete(&self, object_seq: u64) -> Result<()> { + self.inner.delete(object_seq).await } async fn list(&self) -> Result> { @@ -1785,4 +7076,1155 @@ mod tests { self.inner.object_path(object_seq) } } + + /// Drives every outcome of a conditional create through the store, one + /// case per fault, and checks that every row of the *Failure matrix* has one. + mod fault_matrix { + use object_store::layers::mock::{self, MockLayerBuilder, oio}; + use prometheus::Histogram; + + use super::*; + + /// A row of the *Failure matrix*, named by its *Situation* cell. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum Row { + /// create succeeds, or identical retry + Created, + /// transient I/O error, `durable` mode + TransientDurable, + /// an object of an earlier epoch holds the sequence, `durable` mode + EarlierEpochDurable, + /// transient I/O error, `enqueued` mode + TransientEnqueued, + /// transient I/O error, `enqueued` mode after `stop` began + TransientEnqueuedAfterStop, + /// an object of the same or a later epoch holds the sequence, or + /// any conflicting object in the `enqueued` mode + Conflict, + /// encoding or catalog error + CatalogError, + /// create succeeds at the last representable sequence + LastSequence, + /// the start object meets an object of an earlier epoch + StartMovesPast, + /// the start object meets an object of the same or a later epoch, + /// or its create has an unknown outcome + StartFails, + } + + /// Every row of the matrix, in the order the RFC lists them. + const ROWS: [Row; 10] = [ + Row::Created, + Row::TransientDurable, + Row::EarlierEpochDurable, + Row::TransientEnqueued, + Row::TransientEnqueuedAfterStop, + Row::Conflict, + Row::CatalogError, + Row::LastSequence, + Row::StartMovesPast, + Row::StartFails, + ]; + + /// One failure of the object store access and the row it exercises. + struct Case { + fault: &'static str, + row: Row, + run: fn() -> BoxFuture<'static, ()>, + } + + /// Declares [`Counters`] from one sampling expression per field. + macro_rules! counters { + ($($(#[$doc:meta])* $field:ident: $sample:expr,)*) => { + /// The metrics the cases assert on, sampled around a case. + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] + struct Counters { + $($(#[$doc])* $field: u64,)* + } + + impl Counters { + fn sample() -> Self { + Self { $($field: $sample,)* } + } + + /// The metrics recorded since `before`. + fn since(before: Self) -> Self { + let now = Self::sample(); + Self { $($field: now.$field - before.$field,)* } + } + } + }; + } + + counters! { + created_objects: METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL.get(), + /// Samples of the seal-to-durable histogram. + seal_to_durable: METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS.get_sample_count(), + create_failures: METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL.get(), + create_conflicts: METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL.get(), + poisoned: METRIC_OBJECT_STORE_WAL_POISONED_TOTAL.get(), + /// Samples of the object size and entry histograms, one per seal. + sealed_objects: sealed_objects(), + /// The sums of those two histograms. + object_bytes: sum_of(&METRIC_OBJECT_STORE_WAL_OBJECT_BYTES), + object_entries: sum_of(&METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES), + /// Samples of the `durable` acknowledgement histogram. + acknowledged_appends: METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS.get_sample_count(), + /// Objects recovery read footers from, and recoveries timed. + recovered_objects: METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL.get(), + recoveries: METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS.get_sample_count(), + /// Read streams driven to their end. + reads: METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(), + deleted_objects: METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get(), + failed_deletes: METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.get(), + /// No read of the matrix meets a corrupted segment. + skipped_segments: METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(), + stalled_appends: METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get(), + stalled_waits: METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.get_sample_count(), + } + + /// Samples of the two object shape histograms, which a sealed batch + /// feeds together; a batch that fed only one of them fails the case. + fn sealed_objects() -> u64 { + let bytes = METRIC_OBJECT_STORE_WAL_OBJECT_BYTES.get_sample_count(); + let entries = METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES.get_sample_count(); + assert_eq!( + bytes, entries, + "an object was sized but its entries were not counted" + ); + bytes + } + + /// What a histogram of whole numbers has recorded so far. The sums + /// are exact at these magnitudes, so a case states them as counts. + fn sum_of(histogram: &Histogram) -> u64 { + let sum = histogram.get_sample_sum(); + assert_eq!( + sum, + sum.round(), + "a histogram of whole numbers recorded {sum}" + ); + sum as u64 + } + + /// The length of the object a batch of `entries` encodes into; the + /// header and the entry ids have a fixed width, so any will do. + fn object_bytes_of(entries: &[(RegionId, &str)]) -> u64 { + let mut positions = HashMap::new(); + let records = entries + .iter() + .map(|(region_id, data)| { + let position = positions.entry(*region_id).or_insert(0); + *position += 1; + Record { + region_id: *region_id, + entry_id: id(1, *position), + payload: Bytes::from(data.as_bytes().to_vec()), + } + }) + .collect::>(); + encode_object(chain_header(1, 1, None), &records) + .unwrap() + .bytes + .len() as u64 + } + + /// The length of the object of one two-byte entry, which most cases seal. + fn one_entry_object() -> u64 { + object_bytes_of(&[(region(1), "a1")]) + } + + /// The sum and the count of a timing histogram before a forced wait. + fn timer_before(histogram: &Histogram) -> (f64, u64) { + (histogram.get_sample_sum(), histogram.get_sample_count()) + } + + /// One sample of a timing histogram, the wait the case forced on it, + /// and what the case itself took around that wait. + #[derive(Debug, Clone, Copy)] + struct Timed { + recorded: f64, + /// The forced wait as measured, which may be far longer than asked. + held: Duration, + elapsed: Duration, + } + + impl Timed { + fn since( + histogram: &Histogram, + before: (f64, u64), + held: Duration, + elapsed: Duration, + ) -> Self { + assert_eq!( + before.1 + 1, + histogram.get_sample_count(), + "expected the timer to take exactly one sample" + ); + Self { + recorded: histogram.get_sample_sum() - before.0, + held, + elapsed, + } + } + + /// Asserts two observations of one timer: each lies between its own + /// wait and its case's elapsed time, and the long wait outlasted the + /// whole short observation, so a timer that records a constant + /// fails one of them however slow the machine is. + fn assert_pair(name: &str, short: Self, long: Self) { + assert!( + long.held > short.elapsed, + "the {name} timer was observed around a long wait of {:?} that did not outlast the short observation's {:?}", + long.held, + short.elapsed + ); + for observed in [short, long] { + assert!( + observed.recorded >= observed.held.as_secs_f64() + && observed.recorded <= observed.elapsed.as_secs_f64(), + "the {name} timer recorded {}s around a wait of {:?}, expected at most {:?}", + observed.recorded, + observed.held, + observed.elapsed + ); + } + } + } + + fn indexed_gauges() -> (i64, i64) { + ( + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.get(), + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.get(), + ) + } + + /// Asserts that both indexed gauges hold exactly the objects `indexed`, + /// with their sizes as listed. Objects a failed create left behind are + /// listed but not counted. + async fn assert_indexed(io: &dyn WalObjectIo, indexed: &[u64]) { + let bytes = io + .list() + .await + .unwrap() + .into_iter() + .filter(|object| indexed.contains(&object.object_seq)) + .map(|object| object.size) + .sum::(); + assert_eq!( + (indexed.len() as i64, bytes as i64), + indexed_gauges(), + "expected the gauges to hold the objects {indexed:?}" + ); + } + + /// Asserts every append failed with an error `recognise` accepts. + async fn assert_appends_failed( + appends: Vec>>, + recognise: impl Fn(&Error) -> bool, + ) { + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + recognise(unwrap_shared(&error)), + "unexpected error: {error:?}" + ); + } + } + + /// Asserts that a read of the poisoned store fails with the terminal error. + async fn assert_read_is_terminal(store: &ObjectStoreLogStore, region_id: RegionId) { + let error = store + .read(&provider(region_id), 1, None) + .await + .err() + .expect("a read of a poisoned store must fail"); + let terminal = terminal(&store.terminal_error).expect("the store is not poisoned"); + assert!( + matches!(&error, Error::ObjectStoreWal { source, .. } if Arc::ptr_eq(source, &terminal)), + "unexpected error: {error:?}" + ); + } + + /// Polls until `settled` holds. + async fn wait_until(settled: impl Fn() -> bool) { + let deadline = Instant::now() + WAIT; + while !settled() { + assert!(Instant::now() < deadline, "the actor never settled"); + tokio::time::sleep(Duration::from_millis(1)).await; + } + } + + #[tokio::test] + async fn test_store_fault_matrix() { + let _serialized = SERIALIZED.lock().await; + wait_until(|| LIVE_ACTORS.load(Ordering::SeqCst) == 0).await; + let cases = [ + Case { + fault: "a create that succeeds", + row: Row::Created, + run: || Box::pin(create_succeeds()), + }, + Case { + fault: "a create that finds an object with identical content", + row: Row::Created, + run: || Box::pin(create_finds_an_identical_object()), + }, + Case { + fault: "a create whose response is lost", + row: Row::Created, + run: || Box::pin(create_response_is_lost()), + }, + Case { + fault: "a create that fails transiently with later objects created and in flight, `durable` mode", + row: Row::TransientDurable, + run: || Box::pin(create_fails_transiently_durable()), + }, + Case { + fault: "a create that fails transiently, `enqueued` mode", + row: Row::TransientEnqueued, + run: || Box::pin(create_fails_transiently_enqueued()), + }, + Case { + fault: "a create that fails transiently after `stop` began, `enqueued` mode", + row: Row::TransientEnqueuedAfterStop, + run: || Box::pin(create_fails_transiently_after_stop_began()), + }, + Case { + fault: "a create that finds an object of an earlier epoch, `durable` mode", + row: Row::EarlierEpochDurable, + run: || Box::pin(create_finds_an_earlier_epoch()), + }, + Case { + fault: "a create that finds an object of an earlier epoch, `enqueued` mode", + row: Row::Conflict, + run: || Box::pin(create_finds_an_earlier_epoch_enqueued()), + }, + Case { + fault: "a create that finds an object of the same or a later epoch", + row: Row::Conflict, + run: || Box::pin(create_finds_another_writer()), + }, + Case { + fault: "a create that finds an object whose header cannot be read", + row: Row::Conflict, + run: || Box::pin(create_finds_an_unreadable_object()), + }, + Case { + fault: "a create that completes after a conflict poisoned the store", + row: Row::Conflict, + run: || Box::pin(create_completes_after_the_store_was_poisoned()), + }, + Case { + fault: "a created object the catalog rejects", + row: Row::CatalogError, + run: || Box::pin(catalog_rejects_the_created_object()), + }, + Case { + fault: "a create at the last representable object sequence", + row: Row::LastSequence, + run: || Box::pin(create_takes_the_last_object_sequence()), + }, + Case { + fault: "a start object that meets an object of an earlier epoch", + row: Row::StartMovesPast, + run: || Box::pin(start_object_meets_an_earlier_epoch()), + }, + Case { + fault: "a start object that meets an object of the same or a later epoch", + row: Row::StartFails, + run: || Box::pin(start_object_meets_another_writer()), + }, + Case { + fault: "a start object that meets an object with identical content", + row: Row::StartFails, + run: || Box::pin(start_object_meets_an_identical_object()), + }, + Case { + fault: "a start object whose create has an unknown outcome", + row: Row::StartFails, + run: || Box::pin(start_object_create_has_an_unknown_outcome()), + }, + ]; + + common_telemetry::init_default_ut_logging(); + for case in &cases { + common_telemetry::info!("Fault matrix case: {}", case.fault); + (case.run)().await; + } + for row in ROWS { + assert!( + cases.iter().any(|case| case.row == row), + "no case covers the failure matrix row {row:?}" + ); + } + } + + /// The batches are acknowledged in order and the sequence advances. + /// Each create is held, the second past the whole first observation, + /// so both write path timers are observed around disjoint waits. + async fn create_succeeds() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let (region_id, other) = (region(1), region(2)); + let mut seals = Vec::new(); + let mut acknowledgements = Vec::new(); + let mut hold = Duration::from_millis(20); + // The second batch holds entries of two regions, so entries and + // objects differ. + for batch in [ + vec![(region_id, "a1")], + vec![(region_id, "a2"), (region_id, "a3"), (other, "b1")], + ] { + let seal = timer_before(&METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS); + let acknowledgement = timer_before(&METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS); + let started_at = Instant::now(); + let entries = batch + .iter() + .map(|(region_id, data)| entry(&store, *region_id, data)) + .collect(); + let appended = spawn_append_batch(&store, entries); + let (_, release) = next_parked_operation(&mut parked).await; + let parked_at = Instant::now(); + tokio::time::sleep(hold).await; + let held = parked_at.elapsed(); + release.send(true).unwrap(); + timeout(WAIT, appended).await.unwrap().unwrap().unwrap(); + let elapsed = started_at.elapsed(); + let timed = |timer, before| Timed::since(timer, before, held, elapsed); + seals.push(timed( + &METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS, + seal, + )); + acknowledgements.push(timed( + &METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS, + acknowledgement, + )); + hold = elapsed + Duration::from_millis(20); + } + Timed::assert_pair("seal to durable", seals[0], seals[1]); + Timed::assert_pair( + "append acknowledgement", + acknowledgements[0], + acknowledgements[1], + ); + + // Object 1 carries the boundary past the start object. + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!( + expected_entries( + region_id, + &[(id(1, 1), "a1"), (id(2, 1), "a2"), (id(2, 2), "a3")] + ), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[1, 2]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + sealed_objects: 2, + object_bytes: one_entry_object() + + object_bytes_of(&[(region_id, "a2"), (region_id, "a3"), (other, "b1")]), + object_entries: 4, + acknowledged_appends: 2, + reads: 1, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Appends over `io` after `arm` made the create meet the very object + /// it writes, and asserts it was accepted as an identical retry. + async fn assert_identical_retry(io: Arc, arm: impl FnOnce()) { + let store = open_over(io.clone(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + arm(); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!(Some(&id(1, 1)), response.last_entry_ids.get(®ion_id)); + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!(Some(&id(2, 1)), response.last_entry_ids.get(®ion_id)); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(2, 1), "a2")]), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[1, 2]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + acknowledged_appends: 2, + reads: 1, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_finds_an_identical_object() { + let (io, _) = RecordingIo::over(memory_store()); + let armed = io.clone(); + assert_identical_retry(io.clone(), move || { + armed.duplicate_next_put.store(true, Ordering::Relaxed) + }) + .await; + assert!(!io.duplicate_next_put.load(Ordering::Relaxed)); + } + + /// The object store writes the object and loses the response, and the + /// read-back that reconciles the create finds the same bytes. + async fn create_response_is_lost() { + let lose_next = Arc::new(AtomicBool::new(false)); + let losing = lose_next.clone(); + let object_store = memory_store().layer( + MockLayerBuilder::default() + .writer_factory(Arc::new(move |_, _, inner| { + if losing.swap(false, Ordering::SeqCst) { + Box::new(LosingWriter(inner)) + } else { + inner + } + })) + .build() + .unwrap(), + ); + let io = Arc::new(ObjectStoreIo::new(object_store, PREFIX).unwrap()); + let armed = lose_next.clone(); + assert_identical_retry(io, move || armed.store(true, Ordering::SeqCst)).await; + assert!(!lose_next.load(Ordering::SeqCst)); + } + + /// A writer that writes the object and then reports a failure, as a + /// create whose response is lost. + struct LosingWriter(oio::Writer); + + impl oio::Write for LosingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.0.write(buffer).await + } + + async fn close(&mut self) -> mock::Result { + self.0.close().await?; + Err(mock::Error::new(mock::ErrorKind::Unexpected, "lost response").set_temporary()) + } + + async fn abort(&mut self) -> mock::Result<()> { + self.0.abort().await + } + } + + /// Object 2 is created and object 3 in flight when object 1 fails. + async fn create_fails_transiently_durable() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 3).await; + let mut releases = parked_operations(&mut parked, 3).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + releases.remove(&1).unwrap().send(false).unwrap(); + assert_appends_failed(appends, |error| { + matches!(error, Error::WalObjectStore { .. }) + }) + .await; + assert_indexed(io.as_ref(), &[0]).await; + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 3, + object_bytes: 3 * one_entry_object(), + object_entries: 3, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + releases.remove(&3).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 2).await; + let appended = spawn_append_batch(&store, vec![entry(&store, region_id, "a4")]); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(4, object_seq); + release.send(true).unwrap(); + timeout(WAIT, appended).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![2, 3, 4], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(4, 1), "a4")]), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[4]).await; + assert_eq!( + Counters { + created_objects: 3, + seal_to_durable: 3, + create_failures: 1, + sealed_objects: 4, + object_bytes: 4 * one_entry_object(), + object_entries: 4, + acknowledged_appends: 1, + reads: 2, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_fails_transiently_enqueued() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let before = Counters::sample(); + let region_id = region(1); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!(Some(&id(1, 1)), response.last_entry_ids.get(®ion_id)); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + release.send(false).unwrap(); + + // While the repeat is parked nothing is written or indexed. + let (object_seq, repeat) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_eq!(0, store.durable_entry_id(&provider(region_id)).unwrap()); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + repeat.send(true).unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[1]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 2, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_fails_transiently_after_stop_began() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(manual())).await; + let before = Counters::sample(); + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + let seal = { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) + }; + let (_, release) = next_parked_operation(&mut parked).await; + store.begin_stop(); + release.send(false).unwrap(); + timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err(); + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + } + + async fn create_finds_an_earlier_epoch() { + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + put_foreign(&object_store, 1, 0).await; + let before = Counters::sample(); + let region_id = region(1); + let error = append(&store, region_id, "a1").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::StaleWalObject { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + assert_indexed(store.io.as_ref(), &[0]).await; + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_eq!( + Counters { + create_conflicts: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!(Some(&id(2, 1)), response.last_entry_ids.get(®ion_id)); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_indexed(store.io.as_ref(), &[2]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_conflicts: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + acknowledged_appends: 1, + reads: 1, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_finds_an_earlier_epoch_enqueued() { + let object_store = memory_store(); + let store = open(object_store.clone(), &enqueued(eager())).await; + put_foreign(&object_store, 1, 0).await; + let before = Counters::sample(); + let region_id = region(1); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!(Some(&id(1, 1)), response.last_entry_ids.get(®ion_id)); + let error = timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(store.io.as_ref(), &[0]).await; + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + ..Counters::default() + }, + Counters::since(before) + ); + } + + async fn create_finds_another_writer() { + for epoch in [1, 2] { + let object_store = memory_store(); + let (store, io, mut parked) = + open_parking_creates(object_store.clone(), &eager()).await; + put_foreign(&object_store, 1, epoch).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + releases.remove(&1).unwrap().send(true).unwrap(); + let path = io.object_path(1); + assert_appends_failed(appends, |error| { + matches!(error, Error::WalObjectConflict { path: actual, .. } if *actual == path) + }) + .await; + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_conflicts: 1, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + } + + /// The epoch of the object at the sequence cannot be read: the conflict + /// is counted once and the store poisons itself. + async fn create_finds_an_unreadable_object() { + let object_store = memory_store(); + let (io, _) = RecordingIo::over(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + put_foreign(&object_store, 1, 1).await; + let before = Counters::sample(); + let region_id = region(1); + io.damage_next_range_read.store(true, Ordering::Relaxed); + let error = append(&store, region_id, "a1").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::InvalidWalObject { .. }), + "unexpected error: {error:?}" + ); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// The create of object 2 completes after the conflict at object 1 + /// poisoned the store; it is counted and timed all the same. + async fn create_completes_after_the_store_was_poisoned() { + let object_store = memory_store(); + let (store, io, mut parked) = + open_parking_creates(object_store.clone(), &eager()).await; + put_foreign(&object_store, 1, 1).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&1).unwrap().send(true).unwrap(); + assert_appends_failed(appends, |error| { + matches!(error, Error::WalObjectConflict { .. }) + }) + .await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_conflicts: 1, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn catalog_rejects_the_created_object() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + let occupying = FooterEntry { + region_id: region(9), + min_entry_id: id(1, 1), + max_entry_id: id(1, 1), + entry_count: 1, + segment_offset: HEADER_LEN as u64, + segment_len: 1, + segment_crc32: 0, + }; + store + .catalog + .write() + .unwrap() + .insert_object(1, vec![occupying]) + .unwrap(); + releases.remove(&1).unwrap().send(true).unwrap(); + assert_appends_failed(appends, |error| { + matches!(error, Error::CorruptedWalObject { .. }) + }) + .await; + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + // This store indexed neither object, whatever the case put into + // its catalog. + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_takes_the_last_object_sequence() { + let object_store = memory_store(); + let last_object_seq = OBJECT_SEQ_LIMIT - 1; + put_records(&object_store, last_object_seq - 2, &[]).await; + let store = open(object_store, &eager()).await; + let region_id = region(1); + // A poisoned store hands out no entry, so the second is built now. + let second = entry(&store, region_id, "a2"); + let before = Counters::sample(); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!( + Some(&id(last_object_seq, 1)), + response.last_entry_ids.get(®ion_id) + ); + let error = store.append_batch(vec![second]).await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectSequenceExhausted { last_object_seq: actual, .. } + if *actual == last_object_seq), + "unexpected error: {error:?}" + ); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + assert_indexed( + store.io.as_ref(), + &[last_object_seq - 2, last_object_seq - 1, last_object_seq], + ) + .await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + poisoned: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + acknowledged_appends: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Writes object 0 of epoch 1 with one entry of region 1. + async fn put_chain_start(object_store: &ObjectStore) -> ObjectStoreIo { + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let record = Record { + region_id: region(1), + entry_id: id(0, 1), + payload: Bytes::from_static(b"a0"), + }; + put_fixture(&fixtures, 0, &[record]).await; + fixtures + } + + /// Opens a store whose start object meets `late` at sequence 1. + async fn open_meeting( + object_store: &ObjectStore, + late: Bytes, + ) -> Result> { + let io = Arc::new(RacingIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + late: Mutex::new(Some((1, late))), + }); + ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await + } + + async fn start_object_meets_an_earlier_epoch() { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let late = encode_object(chain_header(1, 1, None), &[]).unwrap().bytes; + let before = Counters::sample(); + let store = open_meeting(&object_store, late).await.unwrap(); + assert_eq!(vec![0, 1, 2], object_seqs(&fixtures).await); + assert_eq!( + 3, + decode_header(&fixtures.get(2).await.unwrap()) + .unwrap() + .epoch + ); + assert_indexed(&fixtures, &[0, 2]).await; + assert_eq!( + expected_entries(region(1), &[(id(0, 1), "a0")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!( + Counters { + created_objects: 1, + create_conflicts: 1, + recovered_objects: 1, + recoveries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Asserts a construction failure at the start object: no other + /// sequence was tried, and only the recovery and `create` are counted. + async fn assert_start_fails( + fixtures: &ObjectStoreIo, + before: Counters, + create: Counters, + error: Error, + recognise: impl Fn(&Error) -> bool, + retry_hint: RetryHint, + ) { + assert!(recognise(&error), "unexpected error: {error:?}"); + assert_eq!(retry_hint, error.retry_hint()); + assert_eq!(vec![0, 1], object_seqs(fixtures).await); + assert_indexed(fixtures, &[0]).await; + assert_eq!( + Counters { + recovered_objects: 1, + recoveries: 1, + ..create + }, + Counters::since(before) + ); + } + + async fn start_object_meets_another_writer() { + for epoch in [2, 3] { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let late = encode_object(chain_header(1, epoch, None), &[]) + .unwrap() + .bytes; + let before = Counters::sample(); + let error = open_meeting(&object_store, late).await.unwrap_err(); + let path = fixtures.object_path(1); + assert_start_fails( + &fixtures, + before, + Counters { + create_conflicts: 1, + ..Counters::default() + }, + error, + |error| matches!(error, Error::WalObjectConflict { path: actual, .. } if *actual == path), + RetryHint::NonRetryable, + ) + .await; + } + } + + async fn start_object_meets_an_identical_object() { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let same = encode_object(chain_header(1, 2, Some((0, 1))), &[]) + .unwrap() + .bytes; + let before = Counters::sample(); + let error = open_meeting(&object_store, same).await.unwrap_err(); + assert_start_fails( + &fixtures, + before, + Counters { + created_objects: 1, + ..Counters::default() + }, + error, + |error| matches!(error, Error::UnconfirmedWalEpochStart { epoch: 2, .. }), + RetryHint::Retryable, + ) + .await; + } + + async fn start_object_create_has_an_unknown_outcome() { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let io = Arc::new(LostResponseIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + lose_next: AtomicBool::new(true), + }); + let before = Counters::sample(); + let error = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap_err(); + assert_start_fails( + &fixtures, + before, + Counters { + create_failures: 1, + ..Counters::default() + }, + error, + |error| matches!(error, Error::WalObjectStore { .. }), + RetryHint::Retryable, + ) + .await; + } + } } diff --git a/src/meta-client/Cargo.toml b/src/meta-client/Cargo.toml index 984f434cc51..9863b32bbbf 100644 --- a/src/meta-client/Cargo.toml +++ b/src/meta-client/Cargo.toml @@ -19,7 +19,6 @@ common-options.workspace = true common-telemetry.workspace = true common-time.workspace = true futures.workspace = true -futures-util.workspace = true humantime-serde.workspace = true rand.workspace = true serde.workspace = true @@ -31,10 +30,7 @@ tonic.workspace = true [dev-dependencies] common-event-recorder.workspace = true -datatypes.workspace = true futures.workspace = true meta-srv = { workspace = true, features = ["mock"] } -session.workspace = true -tower.workspace = true tracing = "0.1" tracing-subscriber.workspace = true diff --git a/src/meta-client/src/client/heartbeat.rs b/src/meta-client/src/client/heartbeat.rs index 96fc398556d..6dcc14c13f2 100644 --- a/src/meta-client/src/client/heartbeat.rs +++ b/src/meta-client/src/client/heartbeat.rs @@ -19,7 +19,7 @@ use std::time::Duration; use api::v1::meta::heartbeat_client::HeartbeatClient; use api::v1::meta::{HeartbeatRequest, HeartbeatResponse, RequestHeader, Role}; use common_grpc::channel_manager::ChannelManager; -use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL; +use common_meta::distributed_time_constants::{BASE_HEARTBEAT_INTERVAL, HEARTBEAT_TIMEOUT}; use common_meta::util; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{info, warn}; @@ -249,17 +249,26 @@ impl Inner { })?; let receiver = ReceiverStream::new(receiver); - let mut stream = leader - .heartbeat(receiver) - .await - .map_err(error::Error::from)? - .into_inner(); + let handshake = async { + let mut stream = leader + .heartbeat(receiver) + .await + .map_err(error::Error::from)? + .into_inner(); - let res = stream - .message() + let res = stream + .message() + .await + .map_err(error::Error::from)? + .context(error::CreateHeartbeatStreamSnafu)?; + Ok::<_, error::Error>((stream, res)) + }; + + let (stream, res) = tokio::time::timeout(HEARTBEAT_TIMEOUT, handshake) .await - .map_err(error::Error::from)? - .context(error::CreateHeartbeatStreamSnafu)?; + .with_context(|_| error::HeartbeatHandshakeTimeoutSnafu { + timeout: HEARTBEAT_TIMEOUT, + })??; let config = HeartbeatConfig::from_response(&res); diff --git a/src/meta-client/src/error.rs b/src/meta-client/src/error.rs index 4623cfee33b..e15c83be1d9 100644 --- a/src/meta-client/src/error.rs +++ b/src/meta-client/src/error.rs @@ -86,6 +86,15 @@ pub enum Error { location: Location, }, + #[snafu(display("Heartbeat handshake timeout after {timeout:?}"))] + HeartbeatHandshakeTimeout { + timeout: std::time::Duration, + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: tokio::time::error::Elapsed, + }, + #[snafu(display("Invalid response header"))] InvalidResponseHeader { #[snafu(implicit)] @@ -150,6 +159,7 @@ impl ErrorExt for Error { | Error::MissingQueryContext { .. } | Error::SendHeartbeat { .. } | Error::CreateHeartbeatStream { .. } + | Error::HeartbeatHandshakeTimeout { .. } | Error::CreateChannel { .. } | Error::RetryTimesExceeded { .. } | Error::ConvertMetaConfig { .. } => StatusCode::Internal, @@ -168,6 +178,7 @@ impl ErrorExt for Error { fn retry_hint(&self) -> RetryHint { match self { Error::MetaServer { retry_hint, .. } => *retry_hint, + Error::HeartbeatHandshakeTimeout { .. } => RetryHint::Retryable, Error::InvalidResponseHeader { source, .. } | Error::ConvertMetaRequest { source, .. } | Error::ConvertMetaResponse { source, .. } diff --git a/src/meta-client/tests/heartbeat_handshake_timeout.rs b/src/meta-client/tests/heartbeat_handshake_timeout.rs new file mode 100644 index 00000000000..2ee43fa74d3 --- /dev/null +++ b/src/meta-client/tests/heartbeat_handshake_timeout.rs @@ -0,0 +1,140 @@ +// 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::pin::Pin; +use std::time::Duration; + +use api::v1::meta::heartbeat_server::{Heartbeat, HeartbeatServer}; +use api::v1::meta::{ + AskLeaderRequest, AskLeaderResponse, HeartbeatRequest, HeartbeatResponse, Peer, ResponseHeader, +}; +use common_error::ext::{ErrorExt, RetryHint}; +use common_meta::distributed_time_constants::HEARTBEAT_TIMEOUT; +use futures::Stream; +use meta_client::error::Error; +use meta_client::{MetaClientOptions, MetaClientRef, MetaClientType, create_meta_client}; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::codec::CompressionEncoding; +use tonic::{Request, Response, Status, Streaming}; + +type HeartbeatStream = + Pin> + Send + 'static>>; + +struct MockMetasrv { + addr: String, + send_first_heartbeat_response: bool, +} + +#[async_trait::async_trait] +impl Heartbeat for MockMetasrv { + type HeartbeatStream = HeartbeatStream; + + async fn heartbeat( + &self, + request: Request>, + ) -> Result, Status> { + let send_first_heartbeat_response = self.send_first_heartbeat_response; + let stream = futures::stream::once(async move { + let mut request = request.into_inner(); + request + .message() + .await? + .ok_or_else(|| Status::invalid_argument("missing heartbeat request"))?; + + if send_first_heartbeat_response { + Ok(HeartbeatResponse::default()) + } else { + futures::future::pending::>().await + } + }); + Ok(Response::new(Box::pin(stream))) + } + + async fn ask_leader( + &self, + _request: Request, + ) -> Result, Status> { + Ok(Response::new(AskLeaderResponse { + header: Some(ResponseHeader::default()), + leader: Some(Peer { + id: 1, + addr: self.addr.clone(), + }), + })) + } +} + +async fn create_client(send_first_heartbeat_response: bool) -> MetaClientRef { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let mock = MockMetasrv { + addr: addr.to_string(), + send_first_heartbeat_response, + }; + + tokio::spawn(async move { + let service = HeartbeatServer::new(mock) + .accept_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Zstd) + .send_compressed(CompressionEncoding::Gzip) + .send_compressed(CompressionEncoding::Zstd); + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + + let options = MetaClientOptions { + metasrv_addrs: vec![addr.to_string()], + connect_timeout: Duration::from_millis(500), + ..Default::default() + }; + create_meta_client( + MetaClientType::Datanode { member_id: 1 }, + &options, + None, + None, + ) + .await + .expect("the mock answers leader discovery") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_heartbeat_handshake_succeeds_when_the_server_responds() { + let client = create_client(true).await; + + tokio::time::timeout(HEARTBEAT_TIMEOUT, client.heartbeat()) + .await + .expect("heartbeat handshake must complete before its deadline") + .expect("heartbeat handshake must succeed"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_heartbeat_handshake_times_out_when_the_stream_stays_silent() { + let client = create_client(false).await; + + let result = tokio::time::timeout(2 * HEARTBEAT_TIMEOUT, client.heartbeat()) + .await + .expect("heartbeat handshake must be bounded"); + let error = match result { + Ok(_) => panic!("a silent heartbeat stream must fail the handshake"), + Err(error) => error, + }; + assert!(matches!( + &error, + Error::HeartbeatHandshakeTimeout { timeout, .. } if *timeout == HEARTBEAT_TIMEOUT + )); + assert_eq!(error.retry_hint(), RetryHint::Retryable); +} diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index e1e23581525..5f9f573978c 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -38,7 +38,11 @@ abstraction, election, key encoding, and the DDL manager live in `common-meta`; - **Heartbeat**: `service/heartbeat.rs` receives datanode/frontend heartbeats and runs the handler chain in `handler/`; responses carry mailbox messages (DDL - results, cache invalidation). + results, cache invalidation). Node information is refreshed in memory on each + heartbeat. Datanode statistics collection and flownode remapping share + `handler/utils.rs` for per-node address persistence. Datanode address + updates retain the statistics handler's per-node ordering. Address epochs are + confirmed only after a successful save; cache invalidation remains best effort. - **DDL**: `service/procedure.rs` (`ddl`) is leader-only and hands off to `common-meta`'s `DdlManager`, executed via `common-procedure`. - **Region migration**: `procedure/region_migration/manager.rs` plus per-step diff --git a/src/meta-srv/Cargo.toml b/src/meta-srv/Cargo.toml index e6585505db4..6baca0934e1 100644 --- a/src/meta-srv/Cargo.toml +++ b/src/meta-srv/Cargo.toml @@ -14,16 +14,11 @@ enterprise = ["common-meta/enterprise"] [lints] workspace = true -[target.'cfg(not(target_os = "android"))'.dependencies] -local-ip-address.workspace = true - [dependencies] api.workspace = true async-trait.workspace = true axum.workspace = true -axum-extra.workspace = true axum-macros.workspace = true -bytes.workspace = true chrono.workspace = true clap.workspace = true client.workspace = true @@ -55,18 +50,14 @@ etcd-client.workspace = true futures.workspace = true h2 = "0.3" hostname.workspace = true -http-body-util = "0.1" -humantime.workspace = true humantime-serde.workspace = true hyper-util = { workspace = true, features = ["tokio"] } itertools.workspace = true lazy_static.workspace = true once_cell.workspace = true ordered-float.workspace = true -parking_lot.workspace = true partition.workspace = true prometheus.workspace = true -prost.workspace = true rand.workspace = true regex.workspace = true rskafka.workspace = true @@ -87,7 +78,6 @@ tokio-stream = { workspace = true, features = ["net"] } toml.workspace = true tonic.workspace = true tower.workspace = true -tower-http.workspace = true tracing.workspace = true typetag.workspace = true url = "2.3" diff --git a/src/meta-srv/src/gc/options.rs b/src/meta-srv/src/gc/options.rs index fb43bd6dce1..298fb30bfd8 100644 --- a/src/meta-srv/src/gc/options.rs +++ b/src/meta-srv/src/gc/options.rs @@ -38,7 +38,7 @@ impl Default for SoftDropGcOptions { fn default() -> Self { Self { enable: false, - retention: Duration::from_days(7), + retention: Duration::from_secs(7 * 86400), } } } @@ -242,7 +242,7 @@ mod tests { assert!(!options.experimental_soft_drop.enable); assert_eq!( - Duration::from_days(7), + Duration::from_secs(7 * 86400), options.experimental_soft_drop.retention ); } @@ -254,7 +254,7 @@ mod tests { enable: true, experimental_soft_drop: SoftDropGcOptions { enable: true, - retention: Duration::from_days(1), + retention: Duration::from_secs(86400), }, ..Default::default() }; @@ -269,7 +269,7 @@ mod tests { enable: true, experimental_soft_drop: SoftDropGcOptions { enable: true, - retention: Duration::from_days(1), + retention: Duration::from_secs(86400), }, ..Default::default() }; @@ -284,7 +284,7 @@ mod tests { let options = GcSchedulerOptions { experimental_soft_drop: SoftDropGcOptions { enable: true, - retention: Duration::from_days(1), + retention: Duration::from_secs(86400), }, ..Default::default() }; diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 320f26fae7c..24c1764db2a 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -81,6 +81,7 @@ pub mod publish_heartbeat_handler; pub mod region_lease_handler; pub mod remap_flow_peer_handler; pub mod response_header_handler; +mod utils; #[cfg(test)] pub mod test_utils; @@ -223,18 +224,50 @@ impl Pusher { } #[inline] - pub async fn push(&self, res: HeartbeatResponse) -> Result<()> { - self.sender.send(Ok(res)).await.map_err(|e| { - error::PushMessageSnafu { - err_msg: e.to_string(), - } - .build() - }) + pub fn header(&self) -> ResponseHeader { + self.res_header.clone() + } +} + +/// An owned snapshot of the [`Pusher`] fields the push paths need, so a caller can drop +/// the [`Pushers`] read guard before awaiting the send. It must not hold a [`Pusher`]: +/// dropping one fires the deregister signal, which reads as a closed channel. +struct PusherHandle { + sender: Sender>, + res_header: ResponseHeader, + deregister_signal_receiver: DeregisterSignalReceiver, +} + +impl PusherHandle { + /// Snapshots `pusher`, which callers read under the [`Pushers`] guard. + fn from_pusher(pusher: &Pusher) -> Self { + Self { + sender: pusher.sender.clone(), + res_header: pusher.header(), + deregister_signal_receiver: pusher.deregister_signal_receiver.clone(), + } + } + + /// Sends one mailbox message, consuming the handle. + async fn push(self, mailbox_message: MailboxMessage) -> Result<()> { + self.sender + .send(Ok(HeartbeatResponse { + header: Some(self.res_header), + mailbox_message: Some(mailbox_message), + ..Default::default() + })) + .await + .map_err(|e| { + error::PushMessageSnafu { + err_msg: e.to_string(), + } + .build() + }) } #[inline] - pub fn header(&self) -> ResponseHeader { - self.res_header.clone() + fn deregister_signal_receiver(&self) -> DeregisterSignalReceiver { + self.deregister_signal_receiver.clone() } } @@ -248,41 +281,41 @@ impl Pushers { pusher_id: PusherId, mailbox_message: MailboxMessage, ) -> Result { - let pushers = self.0.read().await; - let pusher = pushers - .get(&pusher_id) - .with_context(|| error::PusherNotFoundSnafu { - pusher_id: pusher_id.to_string(), - })?; + // Releases the guard before the send below, which parks while the target's + // channel is full. `PusherNotFound` is still reported from inside the guard. + let handle = { + let pushers = self.0.read().await; + let pusher = pushers + .get(&pusher_id) + .with_context(|| error::PusherNotFoundSnafu { + pusher_id: pusher_id.to_string(), + })?; + PusherHandle::from_pusher(pusher) + }; - pusher - .push(HeartbeatResponse { - header: Some(pusher.header()), - mailbox_message: Some(mailbox_message), - ..Default::default() - }) - .await?; + let deregister_signal_receiver = handle.deregister_signal_receiver(); + handle.push(mailbox_message).await?; - Ok(pusher.deregister_signal_receiver.clone()) + Ok(deregister_signal_receiver) } async fn broadcast(&self, role: Role, mailbox_message: &MailboxMessage) -> Result<()> { - let pushers = self.0.read().await; - let pushers = pushers - .range(PusherId::role_range(role)) - .map(|(_, value)| value) - .collect::>(); - let mut results = Vec::with_capacity(pushers.len()); + // Collects owned handles, releasing the guard before the sends below. Collecting + // `&Pusher` would hold it for the whole `join_all`. + let targets = { + let pushers = self.0.read().await; + pushers + .range(PusherId::role_range(role)) + .map(|(_, pusher)| PusherHandle::from_pusher(pusher)) + .collect::>() + }; - for pusher in pushers { + let mut results = Vec::with_capacity(targets.len()); + for handle in targets { let mut mailbox_message = mailbox_message.clone(); mailbox_message.id = 0; // one-way message - results.push(pusher.push(HeartbeatResponse { - header: Some(pusher.header()), - mailbox_message: Some(mailbox_message), - ..Default::default() - })) + results.push(handle.push(mailbox_message)); } // Checks the error out of the loop. @@ -923,14 +956,15 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use api::v1::meta::{HeartbeatRequest, MailboxMessage, RequestHeader, Role}; + use api::v1::meta::{HeartbeatRequest, HeartbeatResponse, MailboxMessage, RequestHeader, Role}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; + use futures::future::poll_immediate; use tokio::sync::mpsc; use super::{ HandleControl, HeartbeatAccumulator, HeartbeatHandler, HeartbeatHandlerGroupBuilder, - PusherId, Pushers, + PusherHandle, PusherId, Pushers, }; use crate::error; use crate::handler::collect_stats_handler::CollectStatsHandler; @@ -1396,4 +1430,90 @@ mod tests { let extensions = res.extensions; assert_eq!(extensions.get("custom.key").unwrap(), b"custom-value"); } + + async fn full_pusher( + pushers: &Pushers, + id: PusherId, + ) -> mpsc::Receiver> { + let (tx, rx) = mpsc::channel(1); + tx.try_send(Ok(HeartbeatResponse::default())).unwrap(); + pushers.insert(id, Pusher::new(tx)).await; + // Callers keep `rx` alive: dropping it would make the send fail instead of park. + rx + } + + fn test_mailbox_message(subject: &str) -> MailboxMessage { + MailboxMessage { + id: 0, + subject: subject.to_string(), + timestamp_millis: 123, + ..Default::default() + } + } + + #[tokio::test] + async fn test_pushers_push_does_not_fence_the_registry() { + let pushers = Pushers::default(); + + let full_id = PusherId::new(Role::Datanode, 1); + let _full_rx = full_pusher(&pushers, full_id).await; + + let (healthy_tx, mut healthy_rx) = mpsc::channel(1); + let healthy_id = PusherId::new(Role::Datanode, 2); + pushers.insert(healthy_id, Pusher::new(healthy_tx)).await; + + let msg = test_mailbox_message("fence-test"); + + let mut parked = Box::pin(pushers.push(full_id, msg.clone())); + assert!(poll_immediate(parked.as_mut()).await.is_none()); + + let (writer_tx, _writer_rx) = mpsc::channel(1); + let mut writer = + Box::pin(pushers.insert(PusherId::new(Role::Datanode, 3), Pusher::new(writer_tx))); + assert!(poll_immediate(writer.as_mut()).await.is_some()); + + let mut healthy = Box::pin(pushers.push(healthy_id, msg)); + assert!(poll_immediate(healthy.as_mut()).await.is_some()); + assert!(healthy_rx.try_recv().is_ok()); + } + + #[tokio::test] + async fn test_pushers_broadcast_does_not_fence_the_registry() { + let pushers = Pushers::default(); + + let _full_rx = full_pusher(&pushers, PusherId::new(Role::Datanode, 1)).await; + let (healthy_tx, _healthy_rx) = mpsc::channel(1); + pushers + .insert(PusherId::new(Role::Datanode, 2), Pusher::new(healthy_tx)) + .await; + + let msg = test_mailbox_message("broadcast-fence-test"); + let mut parked = Box::pin(pushers.broadcast(Role::Datanode, &msg)); + assert!(poll_immediate(parked.as_mut()).await.is_none()); + + let (writer_tx, _writer_rx) = mpsc::channel(1); + let mut writer = + Box::pin(pushers.insert(PusherId::new(Role::Flownode, 1), Pusher::new(writer_tx))); + assert!(poll_immediate(writer.as_mut()).await.is_some()); + } + + #[tokio::test] + async fn test_pusher_handle_does_not_fire_the_deregister_signal() { + let (tx, _rx) = mpsc::channel(1); + let pusher = Pusher::new(tx); + let mut signal = pusher.deregister_signal_receiver.clone(); + + let handle = PusherHandle::from_pusher(&pusher); + drop(handle); + + assert!( + tokio::time::timeout(Duration::from_millis(100), signal.changed()) + .await + .is_err() + ); + + // The real `Pusher` still owns the signal. + drop(pusher); + assert!(signal.changed().await.is_ok()); + } } diff --git a/src/meta-srv/src/handler/collect_cluster_info_handler.rs b/src/meta-srv/src/handler/collect_cluster_info_handler.rs index 6bee86ad711..12fee395136 100644 --- a/src/meta-srv/src/handler/collect_cluster_info_handler.rs +++ b/src/meta-srv/src/handler/collect_cluster_info_handler.rs @@ -46,30 +46,21 @@ impl HeartbeatHandler for CollectFrontendClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let frontend_workloads = get_frontend_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Frontend(FrontendStatus { - workloads: frontend_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Frontend(FrontendStatus { + workloads: frontend_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -89,29 +80,20 @@ impl HeartbeatHandler for CollectFlownodeClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let flownode_workloads = get_flownode_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Flownode(FlownodeStatus { - workloads: flownode_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Flownode(FlownodeStatus { + workloads: flownode_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -132,7 +114,7 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { ctx: &mut Context, acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; @@ -147,16 +129,64 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { .count(); let follower_regions = stat.region_stats.len() - leader_regions; - let value = NodeInfo { + update + .save_to_memory( + ctx, + stat.timestamp_millis, + NodeStatus::Datanode(DatanodeStatus { + rcus: stat.rcus, + wcus: stat.wcus, + leader_regions, + follower_regions, + workloads: stat.datanode_workloads.clone(), + }), + ) + .await?; + + Ok(HandleControl::Continue) + } +} + +/// Common fields collected for one heartbeat's in-memory node information +/// update. Persistent routing addresses have a separate lifecycle. +struct NodeInfoUpdate { + key: NodeInfoKey, + peer: Peer, + info: PbNodeInfo, + env_vars: HashMap, +} + +impl NodeInfoUpdate { + fn from_request(request: &HeartbeatRequest) -> Option { + let key = NodeInfoKey::new(request)?; + let peer = request.peer.clone()?; + let info = request.info.clone()?; + let env_vars = EnvVars::from_extensions(&request.extensions) + .inspect_err(|e| { + warn!(e; "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer); + }) + .unwrap_or_default() + .map(|e| e.vars) + .unwrap_or_default(); + Some(Self { + key, peer, - last_activity_ts: stat.timestamp_millis, - status: NodeStatus::Datanode(DatanodeStatus { - rcus: stat.rcus, - wcus: stat.wcus, - leader_regions, - follower_regions, - workloads: stat.datanode_workloads.clone(), - }), + info, + env_vars, + }) + } + + async fn save_to_memory( + self, + ctx: &Context, + last_activity_ts: i64, + status: NodeStatus, + ) -> Result<()> { + let info = self.info; + let value = NodeInfo { + peer: self.peer, + last_activity_ts, + status, version: info.version, git_commit: info.git_commit, start_time_ms: info.start_time_ms, @@ -165,53 +195,88 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { cpu_usage_millicores: info.cpu_usage_millicores, memory_usage_bytes: info.memory_usage_bytes, hostname: info.hostname, - env_vars, + env_vars: self.env_vars, }; - - put_into_memory_store(ctx, key, value).await?; - - Ok(HandleControl::Continue) + let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; + ctx.in_memory + .put(PutRequest { + key: (&self.key).into(), + value, + ..Default::default() + }) + .await + .context(SaveClusterInfoSnafu)?; + Ok(()) } } -fn extract_base_info( - request: &HeartbeatRequest, -) -> Option<(NodeInfoKey, Peer, PbNodeInfo, HashMap)> { - let HeartbeatRequest { peer, info, .. } = request; - let key = NodeInfoKey::new(request)?; - let Some(peer) = &peer else { - return None; - }; - let Some(info) = &info else { - return None; - }; +#[cfg(test)] +mod tests { + use api::v1::meta::RequestHeader; + use common_meta::datanode::Stat; - let env_vars = EnvVars::from_extensions(&request.extensions) - .inspect_err(|e| { - warn!(e; - "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer + use super::*; + use crate::handler::test_utils::TestEnv; + + #[tokio::test] + async fn test_node_info_refreshes_without_epoch_change() { + let mut ctx = TestEnv::new().ctx(); + let handlers: Vec<(Role, Box)> = vec![ + (Role::Frontend, Box::new(CollectFrontendClusterInfoHandler)), + (Role::Datanode, Box::new(CollectDatanodeClusterInfoHandler)), + (Role::Flownode, Box::new(CollectFlownodeClusterInfoHandler)), + ]; + for (role, handler) in handlers { + let mut req = HeartbeatRequest { + header: Some(RequestHeader { + role: role as i32, + ..Default::default() + }), + peer: Some(Peer { + id: 7, + addr: "node:3001".into(), + }), + info: Some(PbNodeInfo { + version: "v1".into(), + hostname: "host".into(), + ..Default::default() + }), + node_epoch: 10, + ..Default::default() + }; + let mut acc = HeartbeatAccumulator { + stat: Some(Stat { + timestamp_millis: 123, + rcus: 5, + wcus: 6, + ..Default::default() + }), + ..Default::default() + }; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let key: Vec = (&NodeInfoKey::new(&req).unwrap()).into(); + let first = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(first.peer, req.peer.clone().unwrap()); + assert_eq!(first.version, "v1"); + assert_eq!(first.hostname, "host"); + req.info.as_mut().unwrap().version = "v2".into(); + req.info.as_mut().unwrap().memory_usage_bytes = 42; + acc.stat.as_mut().unwrap().timestamp_millis = 456; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let second = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(second.version, "v2"); + assert_eq!(second.memory_usage_bytes, 42); + assert_eq!( + serde_json::to_value(&second.status).unwrap(), + serde_json::to_value(&first.status).unwrap() ); - }) - .unwrap_or_default() - .map(|e| e.vars) - .unwrap_or_default(); - - Some((key, peer.clone(), info.clone(), env_vars)) -} - -async fn put_into_memory_store(ctx: &mut Context, key: NodeInfoKey, value: NodeInfo) -> Result<()> { - let key = (&key).into(); - let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; - let put_req = PutRequest { - key, - value, - ..Default::default() - }; - - ctx.in_memory - .put(put_req) - .await - .context(SaveClusterInfoSnafu)?; - - Ok(()) + if role == Role::Datanode { + assert_eq!(second.last_activity_ts, 456); + } + // Node information must not be persisted to the configured backend. + assert!(ctx.kv_backend.get(&key).await.unwrap().is_none()); + } + } } diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index e3ae7ebe915..4b58beaf0fb 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -15,19 +15,19 @@ use std::cmp::Ordering; use std::sync::Arc; -use api::v1::meta::{HeartbeatRequest, Role}; +use api::v1::meta::{HeartbeatRequest, Peer, Role}; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::peer::Peer; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; +use common_telemetry::{error, warn}; use dashmap::DashMap; use snafu::ResultExt; use tokio::sync::Mutex; use crate::error::{self, Result}; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; @@ -40,7 +40,7 @@ struct EpochStats { impl EpochStats { #[inline] fn drain_all(&mut self) -> Vec { - self.stats.drain(..).collect() + std::mem::take(&mut self.stats) } #[inline] @@ -74,6 +74,7 @@ const DEFAULT_FLUSH_STATS_FACTOR: usize = 3; pub struct CollectStatsHandler { stats_cache: DashMap>>, flush_stats_factor: usize, + address_updater: NodeAddressUpdater, } impl Default for CollectStatsHandler { @@ -87,6 +88,7 @@ impl CollectStatsHandler { Self { flush_stats_factor: flush_stats_factor.unwrap_or(DEFAULT_FLUSH_STATS_FACTOR), stats_cache: DashMap::default(), + address_updater: NodeAddressUpdater::default(), } } } @@ -119,6 +121,7 @@ impl HeartbeatHandler for CollectStatsHandler { let key: Vec = key.into(); + let current_epoch = current_stat.node_epoch; let refresh = if let Some(epoch) = epoch_stats.epoch() { match current_stat.node_epoch.cmp(&epoch) { Ordering::Greater => { @@ -146,11 +149,13 @@ impl HeartbeatHandler for CollectStatsHandler { true }; - // Need to refresh the [datanode -> address] mapping - if refresh { - // Safety: `epoch_stats.stats` is not empty - let last = epoch_stats.stats.last().unwrap(); - rewrite_node_address(ctx, last).await; + // Keep address updates and statistics in the original per-node order. + // Address acknowledgement is independent of the statistics epoch so a + // failed save can be retried by the next heartbeat of the same epoch. + if epoch_stats.epoch() == Some(current_epoch) + && let Some(stat) = epoch_stats.stats.last() + { + update_datanode_address(&self.address_updater, ctx, stat).await; } if !refresh && epoch_stats.len() < self.flush_stats_factor { @@ -178,48 +183,26 @@ impl HeartbeatHandler for CollectStatsHandler { } } -async fn rewrite_node_address(ctx: &Context, stat: &Stat) { - let peer = Peer { - id: stat.id, - addr: stat.addr.clone(), - }; - let key = NodeAddressKey::with_datanode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; - - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!( - "Successfully updated datanode `NodeAddressValue`: {:?}", - peer - ); - // broadcast invalidating cache - let cache_idents = stat - .table_ids() - .into_iter() - .map(CacheIdent::TableId) - .collect::>(); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update datanode `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!( - "Failed to serialize datanode `NodeAddressValue`: {:?}", - peer - ); +/// Applies datanode-specific address and table-cache updates. The caller keeps +/// its statistics lock so notifications and statistics retain their ordering. +async fn update_datanode_address(updater: &NodeAddressUpdater, ctx: &Context, stat: &Stat) { + let updated = updater + .update_if_needed((Role::Datanode, stat.id), stat.node_epoch, || async { + let peer = Peer { id: stat.id, addr: stat.addr.clone() }; + let key = NodeAddressKey::with_datanode(stat.id).to_bytes(); + utils::save_node_address(ctx, key, peer).await + }) + .await + .inspect_err(|e| { + error!(e; "Failed to update datanode address, node: {}, address: {}", stat.id, stat.addr); + }); + if let Ok(true) = updated { + let cache_idents = stat + .table_ids() + .into_iter() + .map(CacheIdent::TableId) + .collect::>(); + utils::invalidate_address_caches(ctx, stat.id, &cache_idents).await; } } @@ -553,4 +536,126 @@ mod tests { handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); } } + + #[tokio::test] + async fn test_persistence_failure_retries_but_notification_failure_does_not() { + use crate::handler::collect_stats_handler::CollectStatsHandler; + use crate::handler::test_utils::FailingCacheInvalidator; + use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + use api::v1::meta::HeartbeatRequest; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::store::PutResponse; + let mut ctx = TestEnv::new().ctx(); + let puts = Arc::new(std::sync::Mutex::new(Vec::::new())); + let backend = MockKvBackendBuilder::default() + .put_fn(Arc::new({ + let puts = puts.clone(); + move |req: PutRequest| { + let mut puts = puts.lock().unwrap(); + puts.push(req); + if puts.len() == 1 { + return common_meta::error::UnexpectedSnafu { + err_msg: "injected persistence failure", + } + .fail(); + } + Ok(PutResponse::default()) + } + }) as _) + .build() + .unwrap(); + ctx.leader_cached_kv_backend = + Arc::new(LeaderCachedKvBackend::with_always_leader(Arc::new(backend))); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = CollectStatsHandler::default(); + let mut stat = Stat { + id: 101, + addr: "dn-original".into(), + node_epoch: 7, + ..Default::default() + }; + let req = HeartbeatRequest::default(); + + // A backend error must not interrupt subsequent heartbeat handlers. + assert!(matches!( + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + } + ) + .await + .unwrap(), + HandleControl::Continue + )); + assert_eq!(puts.lock().unwrap().len(), 1); + assert_eq!(invalidator.attempts(), 0); + + let stats_key: Vec = common_meta::datanode::DatanodeStatKey { node_id: 101 }.into(); + assert!(ctx.in_memory.get(&stats_key).await.unwrap().is_some()); + + // Same epoch retries persistence using the current heartbeat's address. + stat.addr = "dn-retry".into(); + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + { + let puts = puts.lock().unwrap(); + assert_eq!(puts.len(), 2); + assert_eq!(puts[1].key, NodeAddressKey::with_datanode(101).to_bytes()); + assert_eq!( + NodeAddressValue::try_from_raw_value(&puts[1].value) + .unwrap() + .peer + .addr, + "dn-retry" + ); + } + assert_eq!(invalidator.attempts(), 1); + + // Notification remains best-effort: its failure neither rewrites nor retries. + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 2); + assert_eq!(invalidator.attempts(), 1); + + stat.node_epoch = 8; + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 3); + assert_eq!(invalidator.attempts(), 2); + } } diff --git a/src/meta-srv/src/handler/failure_handler.rs b/src/meta-srv/src/handler/failure_handler.rs index 5f1448d3ab6..d1d287d13ff 100644 --- a/src/meta-srv/src/handler/failure_handler.rs +++ b/src/meta-srv/src/handler/failure_handler.rs @@ -52,9 +52,10 @@ impl HeartbeatHandler for RegionFailureHandler { return Ok(HandleControl::Continue); }; + // Never blocks: the region lease has already been decided earlier in the + // chain, so a busy supervisor must not withhold the heartbeat response. self.heartbeat_acceptor - .accept(DatanodeHeartbeat::from(stat)) - .await; + .accept(DatanodeHeartbeat::from(stat)); Ok(HandleControl::Continue) } @@ -62,6 +63,8 @@ impl HeartbeatHandler for RegionFailureHandler { #[cfg(test)] mod tests { + use std::time::Duration; + use api::v1::meta::HeartbeatRequest; use common_catalog::consts::default_engine; use common_meta::datanode::{RegionManifestInfo, RegionStat, Stat}; @@ -73,7 +76,45 @@ mod tests { use crate::handler::{HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::builder::MetasrvBuilder; use crate::region::supervisor::tests::new_test_supervisor; - use crate::region::supervisor::{Event, HeartbeatAcceptor}; + use crate::region::supervisor::{DatanodeHeartbeat, Event, HeartbeatAcceptor}; + + fn new_region_stat(region_id: u64) -> RegionStat { + RegionStat { + id: RegionId::from_u64(region_id), + rcus: 0, + wcus: 0, + approximate_bytes: 0, + engine: default_engine().to_string(), + role: RegionRole::Follower, + num_rows: 0, + memtable_size: 0, + manifest_size: 0, + sst_size: 0, + sst_num: 0, + index_size: 0, + region_manifest: RegionManifestInfo::Mito { + manifest_version: 0, + flushed_entry_id: 0, + file_removed_cnt: 0, + }, + data_topic_latest_entry_id: 0, + metadata_topic_latest_entry_id: 0, + written_bytes: 0, + query_cpu_time: 0, + query_scanned_bytes: 0, + min_timestamp: None, + max_timestamp: None, + } + } + + fn new_stat() -> Stat { + Stat { + id: 42, + region_stats: vec![new_region_stat(1), new_region_stat(2), new_region_stat(3)], + timestamp_millis: 1000, + ..Default::default() + } + } #[tokio::test] async fn test_handle_heartbeat() { @@ -85,40 +126,7 @@ mod tests { let metasrv = builder.build().await.unwrap(); let mut ctx = metasrv.new_ctx(); let acc = &mut HeartbeatAccumulator::default(); - fn new_region_stat(region_id: u64) -> RegionStat { - RegionStat { - id: RegionId::from_u64(region_id), - rcus: 0, - wcus: 0, - approximate_bytes: 0, - engine: default_engine().to_string(), - role: RegionRole::Follower, - num_rows: 0, - memtable_size: 0, - manifest_size: 0, - sst_size: 0, - sst_num: 0, - index_size: 0, - region_manifest: RegionManifestInfo::Mito { - manifest_version: 0, - flushed_entry_id: 0, - file_removed_cnt: 0, - }, - data_topic_latest_entry_id: 0, - metadata_topic_latest_entry_id: 0, - written_bytes: 0, - query_cpu_time: 0, - query_scanned_bytes: 0, - min_timestamp: None, - max_timestamp: None, - } - } - acc.stat = Some(Stat { - id: 42, - region_stats: vec![new_region_stat(1), new_region_stat(2), new_region_stat(3)], - timestamp_millis: 1000, - ..Default::default() - }); + acc.stat = Some(new_stat()); handler.handle(req, &mut ctx, acc).await.unwrap(); let (tx, rx) = oneshot::channel(); @@ -126,4 +134,39 @@ mod tests { let detector = rx.await.unwrap(); assert_eq!(detector.iter().collect::>().len(), 3); } + + #[tokio::test] + async fn test_handle_heartbeat_does_not_block_on_full_supervisor_queue() { + // Regression test for #9419: a full queue used to park the handler. + // Keep the supervisor alive, or the queue reports Closed. + let (_supervisor, sender) = new_test_supervisor(); + // Built directly, so the supervisor stays undrained. + let handler = RegionFailureHandler { + heartbeat_acceptor: HeartbeatAcceptor::new(sender.clone()), + }; + + let stat = new_stat(); + let mut queued = 0; + while sender + .try_send(Event::HeartbeatArrived(DatanodeHeartbeat::from(&stat))) + .is_ok() + { + queued += 1; + } + assert_eq!(queued, sender.max_capacity()); + + let builder = MetasrvBuilder::new(); + let metasrv = builder.build().await.unwrap(); + let mut ctx = metasrv.new_ctx(); + let acc = &mut HeartbeatAccumulator::default(); + // Without a stat the handler returns before the acceptor. + acc.stat = Some(stat); + + let result = tokio::time::timeout( + Duration::from_secs(5), + handler.handle(&HeartbeatRequest::default(), &mut ctx, acc), + ) + .await; + assert!(result.is_ok()); + } } diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index 04d63e26e2f..472072aa42b 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -154,9 +154,9 @@ fn align_ts(ts: i64, interval: Duration) -> i64 { impl PersistStatsHandler { /// Creates a new [`PersistStatsHandler`]. pub fn new(inserter: Box, mut persist_interval: Duration) -> Self { - if persist_interval < Duration::from_mins(10) { + if persist_interval < Duration::from_secs(10 * 60) { warn!("persist_interval is less than 10 minutes, set to 10 minutes"); - persist_interval = Duration::from_mins(10); + persist_interval = Duration::from_secs(10 * 60); } Self { diff --git a/src/meta-srv/src/handler/remap_flow_peer_handler.rs b/src/meta-srv/src/handler/remap_flow_peer_handler.rs index f6b3c67a9e5..186cefa19e8 100644 --- a/src/meta-srv/src/handler/remap_flow_peer_handler.rs +++ b/src/meta-srv/src/handler/remap_flow_peer_handler.rs @@ -12,22 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{HeartbeatRequest, Peer, Role}; +use api::v1::meta::{HeartbeatRequest, Role}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; -use dashmap::DashMap; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; +use common_telemetry::error; use crate::Result; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; #[derive(Debug, Default)] pub struct RemapFlowPeerHandler { - /// flow_node_id -> epoch - epoch_cache: DashMap, + address_updater: NodeAddressUpdater, } #[async_trait::async_trait] @@ -45,61 +43,83 @@ impl HeartbeatHandler for RemapFlowPeerHandler { let Some(peer) = req.peer.as_ref() else { return Ok(HandleControl::Continue); }; - - let current_epoch = req.node_epoch; - let flow_node_id = peer.id; - - let refresh = if let Some(mut epoch) = self.epoch_cache.get_mut(&flow_node_id) { - if current_epoch > *epoch.value() { - *epoch.value_mut() = current_epoch; - true - } else { - false - } - } else { - self.epoch_cache.insert(flow_node_id, current_epoch); - true - }; - - if refresh { - rewrite_node_address(ctx, peer).await; + let updated = self + .address_updater + .update_if_needed((Role::Flownode, peer.id), req.node_epoch, || { + let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); + utils::save_node_address(ctx, key, peer.clone()) + }) + .await + .inspect_err(|e| { + // Address errors must not prevent later heartbeat handlers from running. + error!(e; "Failed to update flownode address, peer: {:?}", peer); + }); + if let Ok(true) = updated { + let cache_idents = [CacheIdent::FlowNodeAddressChange(peer.id)]; + utils::invalidate_address_caches(ctx, peer.id, &cache_idents).await; } - Ok(HandleControl::Continue) } } -async fn rewrite_node_address(ctx: &mut Context, peer: &Peer) { - let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; +#[cfg(test)] +mod tests { + use super::*; + use crate::handler::test_utils::TestEnv; + use api::v1::meta::Peer; + use std::sync::Arc; + #[tokio::test] + async fn test_notification_failure_keeps_persisted_epoch() { + use crate::handler::{HeartbeatAccumulator, HeartbeatHandler}; + use api::v1::meta::HeartbeatRequest; - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!("Successfully updated flow `NodeAddressValue`: {:?}", peer); - // broadcast invalidating cache to all frontends - let cache_idents = vec![CacheIdent::FlowNodeAddressChange(peer.id)]; - info!( - "Invalidate flow node cache for new address with cache idents: {:?}", - cache_idents - ); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update flow `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!("Failed to serialize flow `NodeAddressValue`: {:?}", peer); + use crate::handler::test_utils::FailingCacheInvalidator; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + let mut ctx = TestEnv::new().ctx(); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = RemapFlowPeerHandler::default(); + let mut req = HeartbeatRequest { + peer: Some(Peer { + id: 1, + addr: "first".into(), + }), + node_epoch: 1, + ..Default::default() + }; + let key = NodeAddressKey::with_flownode(1).to_bytes(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + req.peer.as_mut().unwrap().addr = "same-epoch".into(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 1); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "first" + ); + req.node_epoch = 2; + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 2); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "same-epoch" + ); } } diff --git a/src/meta-srv/src/handler/test_utils.rs b/src/meta-srv/src/handler/test_utils.rs index d8241c785ea..c2afb3cb14a 100644 --- a/src/meta-srv/src/handler/test_utils.rs +++ b/src/meta-srv/src/handler/test_utils.rs @@ -13,9 +13,13 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; -use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; +use common_meta::cache_invalidator::{ + CacheInvalidator, CacheInvalidatorRef, Context as InvalidationContext, DummyCacheInvalidator, +}; use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL; +use common_meta::instruction::CacheIdent; use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; @@ -97,3 +101,34 @@ impl TestEnv { } } } + +/// Counts failed notification attempts without requiring heartbeat connections. +#[derive(Default)] +pub(crate) struct FailingCacheInvalidator { + attempts: AtomicUsize, +} + +impl FailingCacheInvalidator { + pub(crate) fn attempts(&self) -> usize { + self.attempts.load(Ordering::Relaxed) + } +} + +#[async_trait::async_trait] +impl CacheInvalidator for FailingCacheInvalidator { + async fn invalidate( + &self, + _: &InvalidationContext, + _: &[CacheIdent], + ) -> common_meta::error::Result<()> { + self.attempts.fetch_add(1, Ordering::Relaxed); + common_meta::error::UnexpectedSnafu { + err_msg: "injected notification failure", + } + .fail() + } + + fn invalidate_all(&self) -> common_meta::error::Result<()> { + Ok(()) + } +} diff --git a/src/meta-srv/src/handler/utils.rs b/src/meta-srv/src/handler/utils.rs new file mode 100644 index 00000000000..96624d87306 --- /dev/null +++ b/src/meta-srv/src/handler/utils.rs @@ -0,0 +1,281 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::future::Future; +use std::sync::Arc; + +use api::v1::meta::{Peer, Role}; +use common_meta::instruction::CacheIdent; +use common_meta::key::MetadataValue; +use common_meta::key::node_address::NodeAddressValue; +use common_meta::rpc::store::PutRequest; +use common_telemetry::{error, info}; +use dashmap::DashMap; +use snafu::ResultExt; +use tokio::sync::Mutex; + +use crate::Result; +use crate::error::{InvalidNodeInfoFormatSnafu, KvBackendSnafu}; +use crate::metasrv::Context; + +/// Serializes address changes per node and acknowledges only successful updates. +/// Node information refreshed on every heartbeat has a separate lifecycle. +#[derive(Debug, Default)] +pub(crate) struct NodeAddressUpdater { + states: DashMap<(Role, u64), Arc>>, +} + +#[derive(Debug, Default)] +struct NodeAddressEpochs { + /// Rejects stale heartbeats even after a newer epoch failed to save. + seen_epoch: Option, + /// Suppresses retries only after persistence succeeds. + saved_epoch: Option, +} + +impl NodeAddressUpdater { + /// Saves a new or previously failed epoch. Returns true only when this call + /// saved the address, or false when the heartbeat is stale or already saved. + /// Cache invalidation is separate and does not affect acknowledgement. + pub(crate) async fn update_if_needed( + &self, + node: (Role, u64), + epoch: u64, + save: F, + ) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + // Never retain a DashMap guard across await, even while waiting for a + // different node: its entry may live in the same shard. + let state = { + let entry = self.states.entry(node).or_default(); + Arc::clone(entry.value()) + }; + let mut state = state.lock().await; + if state.seen_epoch.is_some_and(|seen| epoch < seen) || state.saved_epoch == Some(epoch) { + return Ok(false); + } + + // Seeing an epoch rejects stale heartbeats, but only a successful save + // suppresses retries of that same epoch. Cancellation also leaves it retryable. + state.seen_epoch = Some(epoch); + save().await?; + state.saved_epoch = Some(epoch); + Ok(true) + } +} + +/// Persists an address through the leader cache. +pub(crate) async fn save_node_address(ctx: &Context, key: Vec, peer: Peer) -> Result<()> { + let address = NodeAddressValue::new(peer); + let value = address + .try_as_raw_value() + .context(InvalidNodeInfoFormatSnafu)?; + ctx.leader_cached_kv_backend + .put(PutRequest { + key, + value, + prev_kv: false, + }) + .await + .context(KvBackendSnafu)?; + info!("Successfully updated node address: {:?}", address.peer); + Ok(()) +} + +/// Preserves the existing best-effort notification behavior. Notification +/// failures do not undo a persisted address or cause same-epoch rewrites. +pub(crate) async fn invalidate_address_caches( + ctx: &Context, + node_id: u64, + cache_idents: &[CacheIdent], +) { + if let Err(e) = ctx + .cache_invalidator + .invalidate(&Default::default(), cache_idents) + .await + { + error!(e; "Failed to invalidate address caches for node {}: {:?}", node_id, cache_idents); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::sync::Notify; + use tokio::time::timeout; + + use super::*; + + #[tokio::test] + async fn test_failed_save_retries_same_epoch_and_rejects_stale_heartbeats() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + assert!( + updater + .update_if_needed(key, 2, || async { + crate::error::UnexpectedSnafu { + violated: "injected save failure", + } + .fail() + }) + .await + .is_err() + ); + assert!( + !updater + .update_if_needed(key, 1, || async { + panic!("stale heartbeat must not save"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 2, || async { Ok(()) }) + .await + .unwrap() + ); + assert!( + !updater + .update_if_needed(key, 2, || async { + panic!("successful epoch must not save again"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 3, || async { Ok(()) }) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_cancelled_save_can_be_retried() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + let entered = Notify::new(); + { + let first = updater.update_if_needed(key, 0, || async { + entered.notify_one(); + std::future::pending::>().await + }); + tokio::pin!(first); + tokio::select! { + _ = entered.notified() => {}, + _ = &mut first => panic!("save must remain pending"), + } + } + assert!( + timeout( + Duration::from_secs(2), + updater.update_if_needed(key, 0, || async { Ok(()) }) + ) + .await + .unwrap() + .unwrap() + ); + } + + #[test] + fn test_same_node_serialized_other_nodes_progress() { + // A watchdog outside the runtime detects synchronous DashMap deadlocks; + // a Tokio timeout alone cannot fire if the only worker is blocked. + let (progress_tx, progress_rx) = std::sync::mpsc::sync_channel(1); + let release = Arc::new(Notify::new()); + let finish = Arc::new(Notify::new()); + let worker = std::thread::spawn({ + let release = release.clone(); + let finish = finish.clone(); + move || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + .block_on(async move { + let updater = Arc::new(NodeAddressUpdater::default()); + let entered = Arc::new(Notify::new()); + + let first = tokio::spawn({ + let updater = updater.clone(); + let entered = entered.clone(); + let release = release.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 1, || async { + entered.notify_one(); + release.notified().await; + Ok(()) + }) + .await + } + }); + entered.notified().await; + let second = tokio::spawn({ + let updater = updater.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 2, || async { Ok(()) }) + .await + } + }); + tokio::task::yield_now().await; + assert!(!second.is_finished()); + // Exercise shared DashMap shards on a single-thread runtime while the + // first node is suspended in external I/O. + timeout(Duration::from_secs(2), async { + for id in 0_u64..128 { + assert!( + updater + .update_if_needed((Role::Datanode, id + 2), 1, || async { + Ok(()) + }) + .await + .unwrap() + ); + } + }) + .await + .unwrap(); + progress_tx.send(()).unwrap(); + finish.notified().await; + timeout(Duration::from_secs(2), async { + assert!(first.await.unwrap().unwrap()); + assert!(second.await.unwrap().unwrap()); + }) + .await + .unwrap(); + assert!( + !updater + .update_if_needed((Role::Datanode, 1), 1, || async { + panic!("old epoch must not overwrite a newer one"); + }) + .await + .unwrap() + ); + }); + } + }); + let progress = progress_rx.recv_timeout(Duration::from_secs(5)); + release.notify_one(); + finish.notify_one(); + progress.expect("other node updates must progress while one node save is blocked"); + worker.join().unwrap(); + } +} diff --git a/src/meta-srv/src/lib.rs b/src/meta-srv/src/lib.rs index 88e8466be39..948cab240d4 100644 --- a/src/meta-srv/src/lib.rs +++ b/src/meta-srv/src/lib.rs @@ -12,9 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(hash_set_entry)] -#![feature(duration_constructors)] - pub mod bootstrap; pub mod cache_invalidator; pub mod cluster; diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index 65a079bec17..86aee9f8b82 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -122,7 +122,7 @@ impl Default for StatsPersistenceOptions { fn default() -> Self { Self { ttl: Duration::ZERO, - interval: Duration::from_mins(10), + interval: Duration::from_secs(10 * 60), } } } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index b40f06e3cfe..57e7deb3db6 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -92,7 +92,7 @@ use crate::utils::database::DatabaseOperator; use crate::utils::insert_forwarder::InsertForwarder; /// The time window for twcs compaction of the region stats table. -const REGION_STATS_TABLE_TWCS_COMPACTION_TIME_WINDOW: Duration = Duration::from_days(1); +const REGION_STATS_TABLE_TWCS_COMPACTION_TIME_WINDOW: Duration = Duration::from_secs(86400); // TODO(fys): try use derive_builder macro pub struct MetasrvBuilder { @@ -250,6 +250,7 @@ impl MetasrvBuilder { let event_recorder = Arc::new(EventRecorderImpl::with_event_type_filter( Box::new(EventHandlerImpl::new(event_inserter)), options.event_recorder.event_types.clone(), + options.event_recorder.flush_interval, )); let selector = selector.unwrap_or_else(|| Arc::new(LeaseBasedSelector)); diff --git a/src/meta-srv/src/metrics.rs b/src/meta-srv/src/metrics.rs index 1cf32194302..c92e6f31fff 100644 --- a/src/meta-srv/src/metrics.rs +++ b/src/meta-srv/src/metrics.rs @@ -48,6 +48,9 @@ lazy_static! { /// Heartbeat received by metasrv. pub static ref METRIC_META_HEARTBEAT_RECV: IntCounterVec = register_int_counter_vec!("greptime_meta_heartbeat_recv", "heartbeats received by metasrv", &["pusher_key"]).unwrap(); + /// Heartbeats dropped by metasrv. + pub static ref METRIC_META_HEARTBEAT_DROPPED: IntCounterVec = + register_int_counter_vec!("greptime_meta_heartbeat_dropped", "heartbeats dropped by metasrv", &["reason"]).unwrap(); /// The migration execute histogram. pub static ref METRIC_META_REGION_MIGRATION_EXECUTE: HistogramVec = register_histogram_vec!("greptime_meta_region_migration_execute", "meta region migration execute", &["state"]).unwrap(); diff --git a/src/meta-srv/src/mocks.rs b/src/meta-srv/src/mocks.rs index 1a6fa306743..1afc5dbd334 100644 --- a/src/meta-srv/src/mocks.rs +++ b/src/meta-srv/src/mocks.rs @@ -21,6 +21,7 @@ use api::v1::meta::heartbeat_server::HeartbeatServer; use api::v1::meta::procedure_service_server::ProcedureServiceServer; use api::v1::meta::store_server::StoreServer; use client::client_manager::NodeClients; +use common_base::Plugins; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::key::TableMetadataManager; use common_meta::kv_backend::etcd::EtcdStore; @@ -95,6 +96,7 @@ pub async fn mock( datanode_clients, in_memory, None, + Plugins::default(), ) .await } @@ -114,6 +116,30 @@ pub async fn mock_with_client_channel_config( datanode_clients, in_memory, Some(client_channel_config), + Plugins::default(), + ) + .await +} + +/// Builds a mock metasrv with the given [`Plugins`]. +/// +/// Same as [`mock`], except that the plugins are passed to the metasrv builder. +pub async fn mock_with_plugins( + opts: MetasrvOptions, + kv_backend: KvBackendRef, + selector: Option, + datanode_clients: Option>, + in_memory: Option, + plugins: Plugins, +) -> MockInfo { + mock_inner( + opts, + kv_backend, + selector, + datanode_clients, + in_memory, + None, + plugins, ) .await } @@ -125,6 +151,7 @@ async fn mock_inner( datanode_clients: Option>, in_memory: Option, client_channel_config: Option, + plugins: Plugins, ) -> MockInfo { let server_addr = opts.grpc.server_addr.clone(); let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); @@ -152,6 +179,8 @@ async fn mock_inner( None => builder, }; + let builder = builder.plugins(plugins); + let metasrv = builder.build().await.unwrap(); metasrv.try_start().await.unwrap(); diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index f42d8aacd0c..008c0ce05bb 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -107,7 +107,7 @@ pub struct PersistentContext { } fn default_timeout() -> Duration { - Duration::from_mins(2) + Duration::from_secs(2 * 60) } impl PersistentContext { diff --git a/src/meta-srv/src/procedure/repartition/utils.rs b/src/meta-srv/src/procedure/repartition/utils.rs index 50dbc57f9a8..cf6ea726f34 100644 --- a/src/meta-srv/src/procedure/repartition/utils.rs +++ b/src/meta-srv/src/procedure/repartition/utils.rs @@ -82,7 +82,7 @@ pub fn merge_and_validate_region_wal_options( table_id: TableId, ) -> Result { // Doesn't allow overwriting existing WAL options. - for (region_number, _) in new_region_wal_options.iter() { + for region_number in new_region_wal_options.keys() { if region_wal_options.contains_key(region_number) { return error::UnexpectedSnafu { violated: format!( diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index c30519da984..49964de24b7 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -336,6 +336,9 @@ pub mod mock { }), ) .unwrap(); + // The mock transport is a single in-memory duplex connection, so both lanes + // must resolve through this one pre-wired manager. + #[allow(deprecated)] Client::with_manager_and_urls(channel_manager, vec![datanode.addr.clone()]) } } diff --git a/src/meta-srv/src/procedure/wal_prune/manager.rs b/src/meta-srv/src/procedure/wal_prune/manager.rs index e398be93e45..4de21086906 100644 --- a/src/meta-srv/src/procedure/wal_prune/manager.rs +++ b/src/meta-srv/src/procedure/wal_prune/manager.rs @@ -13,7 +13,6 @@ // limitations under the License. use std::collections::HashSet; -use std::collections::hash_set::Entry; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, RwLock}; @@ -46,15 +45,13 @@ impl WalPruneProcedureTracker { /// consume acquire a semaphore permit for the given topic name. pub fn insert_running_procedure(&self, topic_name: String) -> Option { let mut running_procedures = self.running_procedures.write().unwrap(); - match running_procedures.entry(topic_name.clone()) { - Entry::Occupied(_) => None, - Entry::Vacant(entry) => { - entry.insert(); - Some(WalPruneProcedureGuard { - topic_name, - running_procedures: self.running_procedures.clone(), - }) - } + if running_procedures.insert(topic_name.clone()) { + Some(WalPruneProcedureGuard { + topic_name, + running_procedures: self.running_procedures.clone(), + }) + } else { + None } } diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 56a9fe4ee57..99c68ac23c8 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -35,6 +35,7 @@ use common_time::util::current_time_millis; use futures::{StreamExt, TryStreamExt}; use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; +use tokio::sync::mpsc::error::TrySendError; use tokio::sync::mpsc::{Receiver, Sender}; use tokio::sync::oneshot; use tokio::time::{MissedTickBehavior, interval, interval_at}; @@ -43,6 +44,7 @@ use crate::discovery::utils::accept_ingest_workload; use crate::error::{self, Result}; use crate::failure_detector::PhiAccrualFailureDetectorOptions; use crate::metasrv::{RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorRef}; +use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::manager::{ RegionMigrationManagerRef, RegionMigrationTriggerReason, SubmitRegionMigrationTaskResult, }; @@ -356,10 +358,29 @@ impl HeartbeatAcceptor { Self { sender } } - /// Accepts heartbeats from datanodes. - pub(crate) async fn accept(&self, heartbeat: DatanodeHeartbeat) { - if let Err(err) = self.sender.send(Event::HeartbeatArrived(heartbeat)).await { - error!(err; "RegionSupervisor has stop receiving heartbeat."); + /// Accepts heartbeats from datanodes without waiting for queue capacity: a full queue + /// drops the heartbeat, counted in [`METRIC_META_HEARTBEAT_DROPPED`]. A dropped + /// heartbeat costs the failure detector one sample. + pub(crate) fn accept(&self, heartbeat: DatanodeHeartbeat) { + let datanode_id = heartbeat.datanode_id; + match self.sender.try_send(Event::HeartbeatArrived(heartbeat)) { + Ok(()) => {} + Err(TrySendError::Full(_)) => { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["full"]) + .inc(); + warn!( + "Dropping heartbeat because the region supervisor event queue is full, datanode_id: {}, queue_capacity: {}", + datanode_id, + self.sender.max_capacity() + ); + } + Err(TrySendError::Closed(err)) => { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["closed"]) + .inc(); + error!(err; "RegionSupervisor has stop receiving heartbeat."); + } } } } @@ -899,14 +920,15 @@ pub(crate) mod tests { use tokio::time::sleep; use super::RegionSupervisorSelector; + use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::RegionMigrationTriggerReason; use crate::procedure::region_migration::manager::{ RegionMigrationManager, SubmitRegionMigrationTaskResult, }; use crate::procedure::region_migration::test_util::TestingEnv; use crate::region::supervisor::{ - DatanodeHeartbeat, Event, RegionFailureDetectorControl, RegionSupervisor, - RegionSupervisorTicker, + DatanodeHeartbeat, Event, HeartbeatAcceptor, RegionFailureDetectorControl, + RegionSupervisor, RegionSupervisorTicker, }; use crate::selector::test_utils::{RandomNodeSelector, new_test_selector_context}; @@ -940,6 +962,14 @@ pub(crate) mod tests { ) } + fn filler_heartbeat() -> DatanodeHeartbeat { + DatanodeHeartbeat { + datanode_id: 1, + regions: vec![], + timestamp: 0, + } + } + #[tokio::test] async fn test_heartbeat() { let (mut supervisor, sender) = new_test_supervisor(); @@ -1372,4 +1402,39 @@ pub(crate) mod tests { assert!(supervisor.failure_detector.contains(&detecting_region)); assert!(supervisor.failover_counts.contains_key(&detecting_region)); } + + #[tokio::test] + async fn test_heartbeat_accept_does_not_block_on_full_queue() { + // Regression test for #9419: a full queue used to park the caller. + // The supervisor is never run, so the queue stays full. + let (mut supervisor, sender) = new_test_supervisor(); + + let mut queued = 0; + while sender + .try_send(Event::HeartbeatArrived(filler_heartbeat())) + .is_ok() + { + queued += 1; + } + assert_eq!(queued, sender.max_capacity()); + assert_eq!(sender.capacity(), 0); + + let acceptor = HeartbeatAcceptor::new(sender.clone()); + let dropped = || { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["full"]) + .get() + }; + let dropped_before = dropped(); + // The full queue drops the heartbeat and counts it. + acceptor.accept(filler_heartbeat()); + assert_eq!(dropped(), dropped_before + 1); + assert_eq!(sender.capacity(), 0); + + // Free one slot: the next heartbeat is accepted. + supervisor.receiver.try_recv().unwrap(); + acceptor.accept(filler_heartbeat()); + assert_eq!(dropped(), dropped_before + 1); + assert_eq!(sender.capacity(), 0); + } } diff --git a/src/meta-srv/src/service/heartbeat.rs b/src/meta-srv/src/service/heartbeat.rs index 5cdc5b57894..782daa289de 100644 --- a/src/meta-srv/src/service/heartbeat.rs +++ b/src/meta-srv/src/service/heartbeat.rs @@ -419,7 +419,6 @@ mod tests { use std::sync::Arc; use api::v1::meta::heartbeat_server::Heartbeat; - use api::v1::meta::*; use common_meta::kv_backend::memory::MemoryKvBackend; use common_telemetry::tracing_context::W3cTrace; use servers::grpc::GrpcOptions; diff --git a/src/meta-srv/src/utils/database.rs b/src/meta-srv/src/utils/database.rs index 976be8a85b6..796d63009d8 100644 --- a/src/meta-srv/src/utils/database.rs +++ b/src/meta-srv/src/utils/database.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use client::error::{ExternalSnafu, Result as ClientResult}; use client::{Client, Database, Output}; use common_error::ext::BoxedError; +use common_grpc::channel_manager::ChannelManager; use common_meta::peer::PeerDiscoveryRef; use common_telemetry::{debug, warn}; use snafu::{ResultExt, ensure}; @@ -43,6 +44,9 @@ impl<'a> DatabaseContext<'a> { } /// A cached frontend database operator used by metasrv. +/// +/// Its cached clients use independent query and control channel-manager pools; +/// the legacy single-manager client constructors intentionally remain shared. pub struct DatabaseOperator { peer_discovery: PeerDiscoveryRef, client: RwLock>, @@ -110,7 +114,11 @@ impl DatabaseOperator { debug!("Available frontend addresses: {:?}", urls); - Ok(Client::with_urls(urls)) + Ok(Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + urls, + )) } async fn maybe_init_client(&self) -> ClientResult { @@ -147,3 +155,62 @@ fn should_reset_client(result: &client::error::Result) -> bool { .map(|err| err.is_connection_error()) .unwrap_or(false) } + +#[cfg(test)] +mod tests { + use api::v1::meta::heartbeat_request::NodeWorkloads; + use common_meta::cluster::{FrontendStatus, NodeInfo, NodeStatus}; + use common_meta::peer::{Peer, PeerDiscovery}; + + use super::*; + + struct TestPeerDiscovery; + + #[async_trait::async_trait] + impl PeerDiscovery for TestPeerDiscovery { + async fn active_frontends(&self) -> common_meta::error::Result> { + Ok(vec![NodeInfo { + peer: Peer::new(1, "127.0.0.1:3001".to_string()), + last_activity_ts: 0, + status: NodeStatus::Frontend(FrontendStatus::default()), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars: Default::default(), + }]) + } + + async fn active_datanodes( + &self, + _filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> common_meta::error::Result> { + unreachable!() + } + + async fn active_flownodes( + &self, + _filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> common_meta::error::Result> { + unreachable!() + } + } + + #[tokio::test] + async fn test_build_client_uses_isolated_reused_channel_pools() { + let operator = DatabaseOperator::new(Arc::new(TestPeerDiscovery)); + let client = operator.build_client().await.unwrap(); + + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_eq!((1, 0), client.channel_pool_sizes()); + + client.find_channel().unwrap(); + client.find_channel().unwrap(); + assert_eq!((1, 1), client.channel_pool_sizes()); + } +} diff --git a/src/meta-srv/src/utils/postgres.rs b/src/meta-srv/src/utils/postgres.rs index ec46c076e26..148a654e2cd 100644 --- a/src/meta-srv/src/utils/postgres.rs +++ b/src/meta-srv/src/utils/postgres.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(unix)] +use std::str::FromStr; + use common_error::ext::BoxedError; use common_meta::election::ElectionRef; use common_meta::election::rds::postgres::{ElectionPgClient, PgElection}; @@ -20,10 +23,13 @@ use common_meta::kv_backend::rds::PgStore; use common_meta::kv_backend::rds::postgres::{ TlsMode as PgTlsMode, TlsOption as PgTlsOption, create_postgres_tls_connector, }; +use common_telemetry::warn; use deadpool_postgres::{Config, Runtime}; use servers::tls::TlsOption; use snafu::{OptionExt, ResultExt}; use tokio_postgres::NoTls; +#[cfg(unix)] +use tokio_postgres::config::Host; use crate::error::{self, Result}; @@ -60,23 +66,57 @@ pub async fn create_postgres_pool( })?; cfg.url = Some(postgres_url.clone()); - let pool = if let Some(tls_config) = tls_config { - let pg_tls_config = convert_tls_option(&tls_config); - let tls_connector = - create_postgres_tls_connector(&pg_tls_config).map_err(|e| error::Error::Other { - source: BoxedError::new(e), - location: snafu::Location::new(file!(), line!(), 0), - })?; - cfg.create_pool(Some(Runtime::Tokio1), tls_connector) - .context(error::CreatePostgresPoolSnafu)? - } else { - cfg.create_pool(Some(Runtime::Tokio1), NoTls) - .context(error::CreatePostgresPoolSnafu)? + let is_unix_socket = is_unix_socket_url(postgres_url); + if is_unix_socket + && matches!(tls_config.as_ref(), Some(t) if t.mode != servers::tls::TlsMode::Disable) + { + warn!( + "TLS is not supported for Unix domain socket PostgreSQL connections, falling back to NoTls" + ); + } + + let pool = match tls_config { + Some(tls_config) + if tls_config.mode != servers::tls::TlsMode::Disable && !is_unix_socket => + { + let pg_tls_config = convert_tls_option(&tls_config); + let tls_connector = + create_postgres_tls_connector(&pg_tls_config).map_err(|e| error::Error::Other { + source: BoxedError::new(e), + location: snafu::Location::new(file!(), line!(), 0), + })?; + cfg.create_pool(Some(Runtime::Tokio1), tls_connector) + .context(error::CreatePostgresPoolSnafu)? + } + _ => cfg + .create_pool(Some(Runtime::Tokio1), NoTls) + .context(error::CreatePostgresPoolSnafu)?, }; Ok(pool) } +#[cfg(unix)] +fn is_unix_socket_url(url: &str) -> bool { + let Ok(cfg) = tokio_postgres::Config::from_str(url) else { + return false; + }; + // tokio-postgres dials `hostaddr` over TCP even when `host` is a socket path, + // so treat the config as a socket only when every host is a Unix socket and + // no `hostaddr` is set. + cfg.get_hostaddrs().is_empty() + && !cfg.get_hosts().is_empty() + && cfg + .get_hosts() + .iter() + .all(|host| matches!(host, Host::Unix(_))) +} + +#[cfg(not(unix))] +fn is_unix_socket_url(_: &str) -> bool { + false +} + /// Builds a Postgres-backed metadata [`KvBackendRef`]. /// /// * `store_addrs` - Postgres connection URLs; only the first address is used. @@ -150,3 +190,59 @@ pub async fn build_postgres_election( .await .context(error::KvBackendSnafu) } + +#[cfg(test)] +mod tests { + use super::is_unix_socket_url; + + #[test] + fn detects_postgres_unix_socket_url() { + #[cfg(unix)] + { + // libpq keyword-value form (issue #7734) + assert!(is_unix_socket_url( + "host=/var/run/postgresql dbname=greptime user=greptime password=secret" + )); + // standard postgres URL with percent-encoded unix socket directory + assert!(is_unix_socket_url( + "postgresql://user:pw@%2Fvar%2Frun%2Fpostgresql/mydb" + )); + assert!(is_unix_socket_url( + "postgresql://user@%2Fvar%2Frun%2Fpostgresql/db" + )); + // postgres URL with socket dir in query param + assert!(is_unix_socket_url( + "postgresql:///mydb?host=%2Fvar%2Frun%2Fpostgresql" + )); + assert!(is_unix_socket_url( + "postgresql://user:secret@/mydb?host=%2Fvar%2Frun%2Fpostgresql" + )); + + // TCP URLs should not be classified as unix socket + assert!(!is_unix_socket_url( + "postgresql://user:pw@localhost:5432/mydb" + )); + assert!(!is_unix_socket_url("postgresql://user@localhost/db")); + assert!(!is_unix_socket_url( + "host=127.0.0.1 port=5432 dbname=greptime user=greptime password=secret" + )); + // mixed socket and TCP hosts must not be treated as a socket-only config: + // a Require/VerifyFull TLS config would otherwise end up in plaintext over TCP + assert!(!is_unix_socket_url( + "host=/var/run/postgresql,db.example.com" + )); + // tokio-postgres dials `hostaddr` over TCP even when `host` is a socket path + assert!(!is_unix_socket_url("host=/tmp hostaddr=10.0.0.5")); + } + + #[cfg(not(unix))] + { + assert!(!is_unix_socket_url( + "host=/var/run/postgresql dbname=greptime user=greptime password=secret" + )); + assert!(!is_unix_socket_url( + "postgresql://user:pw@localhost:5432/mydb" + )); + } + } +} diff --git a/src/metric-engine/Cargo.toml b/src/metric-engine/Cargo.toml index 0005a4dc543..92000549832 100644 --- a/src/metric-engine/Cargo.toml +++ b/src/metric-engine/Cargo.toml @@ -14,7 +14,7 @@ async-stream.workspace = true async-trait.workspace = true base64.workspace = true bytes.workspace = true -fxhash = "0.2" +chrono.workspace = true common-base.workspace = true common-error.workspace = true common-grpc.workspace = true @@ -24,10 +24,10 @@ common-recordbatch.workspace = true common-runtime.workspace = true common-telemetry.workspace = true common-time.workspace = true -chrono.workspace = true datafusion.workspace = true datatypes.workspace = true futures-util.workspace = true +fxhash = "0.2" humantime-serde.workspace = true itertools.workspace = true lazy_static = "1.4" @@ -45,18 +45,14 @@ tokio.workspace = true tracing.workspace = true [dev-dependencies] +common-function = { workspace = true } common-meta = { workspace = true, features = ["testing"] } -common-test-util.workspace = true -mito2 = { workspace = true, features = ["test"] } common-wal = { workspace = true } criterion = { workspace = true, features = ["async", "async_tokio"] } +mito2 = { workspace = true, features = ["test"] } mur3 = "0.1" -common-function = { workspace = true } partition = { workspace = true } [[bench]] name = "bench_tsid_generator" harness = false - -[package.metadata.cargo-udeps.ignore] -normal = ["aquamarine"] diff --git a/src/metric-engine/src/data_region.rs b/src/metric-engine/src/data_region.rs index 55b9eaaf683..a11b3e57b54 100644 --- a/src/metric-engine/src/data_region.rs +++ b/src/metric-engine/src/data_region.rs @@ -227,6 +227,7 @@ impl DataRegion { | AlterKind::UnsetRegionOptions { keys: _ } | AlterKind::SetIndexes { options: _ } | AlterKind::UnsetIndexes { options: _ } + | AlterKind::SetJsonSettings { .. } | AlterKind::SyncColumns { column_metadatas: _, } => { diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 18d6ff6f197..ec180839526 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -614,6 +614,37 @@ mod test { TestEnv, build_rows, create_logical_region_request, row_schema_with_tags, }; + #[tokio::test] + async fn test_build_series_index_forwarding() { + use api::v1::region::build_index_request; + use store_api::region_request::RegionBuildIndexRequest; + + let env = TestEnv::new().await; + env.init_metric_region().await; + let request = || { + RegionRequest::BuildIndex(RegionBuildIndexRequest { + options: Some(build_index_request::Options::SeriesIndex(Default::default())), + }) + }; + // Default Mito config disables series indexes: reaching that error proves forwarding. + let error = env + .metric() + .handle_request(env.default_physical_region_id(), request()) + .await + .unwrap_err(); + let error = format!("{error:?}"); + assert!(error.contains("series index is disabled"), "{error}"); + let error = env + .metric() + .handle_request(env.default_logical_region_id(), request()) + .await + .unwrap_err(); + assert_eq!( + common_error::status_code::StatusCode::Unsupported, + error.status_code() + ); + } + #[tokio::test] async fn test_discard_unflushed_data_only() { let env = TestEnv::new().await; diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index cb43a71406d..dd20f34adf8 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -18,10 +18,9 @@ use api::v1::{ArrowIpc, SemanticType}; use bytes::Bytes; use common_grpc::flight::{FlightEncoder, FlightMessage}; use datatypes::arrow::record_batch::RecordBatch; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{OptionExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; -use store_api::region_engine::RegionEngine; use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest}; use store_api::storage::RegionId; @@ -74,9 +73,9 @@ impl MetricEngineInner { region_id: RegionId, mut request: RegionBulkInsertsRequest, ) -> Result { - // Simply set the aligned schema to the data region schema version to avoid filling missing columns - // because that schema should be constant and callers have ensured request has the same schema. - request.aligned_schema_version = Some(self.physical_schema_version(region_id).await?); + // Other logical tables can add fields to the shared physical schema. + // Let Mito fill fields absent from this batch before writing it. + request.aligned_schema_version = None; self.data_region .write_data(region_id, RegionRequest::BulkInserts(request)) .await @@ -121,7 +120,6 @@ impl MetricEngineInner { let (schema, data_header, payload) = record_batch_to_ipc(&modified_batch)?; let partition_expr_version = request.partition_expr_version; - let aligned_schema_version = Some(self.physical_schema_version(data_region_id).await?); let request = RegionBulkInsertsRequest { skip_wal: request.skip_wal, @@ -133,22 +131,13 @@ impl MetricEngineInner { payload, }, partition_expr_version, - aligned_schema_version, + aligned_schema_version: None, }; self.data_region .write_data(data_region_id, RegionRequest::BulkInserts(request)) .await } - async fn physical_schema_version(&self, region_id: RegionId) -> Result { - Ok(self - .mito - .get_metadata(region_id) - .await - .context(error::MitoReadOperationSnafu)? - .schema_version) - } - fn resolve_tag_columns_from_metadata( &self, logical_region_id: RegionId, @@ -365,17 +354,23 @@ 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; + check_bulk_insert_physical_region_passthrough(false, false).await; + check_bulk_insert_physical_region_passthrough(true, false).await; } - async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool) { + #[tokio::test] + async fn test_bulk_insert_physical_region_close_recovery() { + check_bulk_insert_physical_region_passthrough(false, true).await; + check_bulk_insert_physical_region_passthrough(true, true).await; + } + + async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool, close: bool) { // Use flat format so that BulkMemtable is used (supports write_bulk). let mito_config = MitoConfig { default_flat_format: true, ..Default::default() }; - let env = TestEnv::with_mito_config("", mito_config, Default::default()).await; + let mut 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(); @@ -425,20 +420,23 @@ mod tests { 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(); + let reopened = if close { + env.metric() + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + // Recreating the wrapper prevents its metadata cache from hiding data loss. + MetricEngine::try_new(env.mito(), Default::default()).unwrap() + } else { + // Stopping workers without Close exercises recovery from WAL alone. + env.reopen_engine(Default::default()).await; + env.metric() + }; reopened.inner.flush_task.stop().await.unwrap(); reopened .handle_request( @@ -457,6 +455,11 @@ mod tests { ) .await .unwrap(); + let stat = env + .mito() + .region_statistic(crate::utils::to_data_region_id(physical_region_id)) + .unwrap(); + assert_eq!(stat.sst_num > 0, skip_wal && close); let stream = reopened .scan_to_stream(logical_region_id, ScanRequest::default()) .await @@ -464,7 +467,7 @@ mod tests { let batches = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!( batches.iter().map(|b| b.num_rows()).sum::(), - if skip_wal { 0 } else { 6 }, + if skip_wal && !close { 0 } else { 6 }, ); } diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 56d4af2284f..b709d99d847 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -791,7 +791,9 @@ 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::arrow::array::{ + Float64Array, TimestampMicrosecondArray, TimestampMillisecondArray, + }; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use datatypes::value::Value as PartitionValue; @@ -932,14 +934,22 @@ mod tests { #[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; + check_put_skip_wal_batch_recovery("sparse", false, false).await; + check_put_skip_wal_batch_recovery("sparse", true, false).await; + check_put_skip_wal_batch_recovery("dense", false, false).await; + check_put_skip_wal_batch_recovery("dense", true, false).await; } - async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool) { - let env = TestEnv::new().await; + #[tokio::test] + async fn test_put_skip_wal_batch_close_recovery() { + check_put_skip_wal_batch_recovery("sparse", false, true).await; + check_put_skip_wal_batch_recovery("sparse", true, true).await; + check_put_skip_wal_batch_recovery("dense", false, true).await; + check_put_skip_wal_batch_recovery("dense", true, true).await; + } + + async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool, close: bool) { + let mut env = TestEnv::new().await; let engine = env.metric(); engine.inner.flush_task.stop().await.unwrap(); let physical_region_id = env.default_physical_region_id(); @@ -1007,18 +1017,21 @@ mod tests { 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(); + let reopened = if close { + engine + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + // Recreate the wrapper as well, discarding its metadata cache. + MetricEngine::try_new(env.mito(), Default::default()).unwrap() + } else { + // Bypass Close to retain coverage of rows lost without WAL protection. + env.reopen_engine(Default::default()).await; + env.metric() + }; reopened.inner.flush_task.stop().await.unwrap(); reopened .handle_request( @@ -1040,12 +1053,17 @@ mod tests { ) .await .unwrap(); + let stat = env + .mito() + .region_statistic(crate::utils::to_data_region_id(physical_region_id)) + .unwrap(); + assert_eq!(stat.sst_num > 0, skip_wal && close); 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 { + let expected = if skip_wal && !close { vec![] } else { vec![(0, 30.0), (1, 10.0), (2, 20.0), (3, 30.0)] @@ -1053,7 +1071,7 @@ mod tests { assert_eq!( scan_timestamp_values(&reopened, logical_region_id).await, expected, - "encoding={encoding}, skip_wal={skip_wal}" + "encoding={encoding}, skip_wal={skip_wal}, close={close}" ); } @@ -1116,6 +1134,91 @@ mod tests { expr.as_json_str().unwrap() } + #[tokio::test] + async fn test_put_and_scan_microsecond_physical_region() { + let env = TestEnv::new().await; + let engine = env.metric(); + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + env.create_physical_region_with_ts_type( + physical_region_id, + &TestEnv::default_table_dir(), + vec![], + ConcreteDataType::timestamp_microsecond_datatype(), + ) + .await; + + // A logical region with a matching microsecond time index is accepted. + let region_create_request = test_util::create_logical_region_request_with_ts_type( + &["job"], + physical_region_id, + &table_dir("test", logical_region_id.table_id()), + ConcreteDataType::timestamp_microsecond_datatype(), + ); + engine + .handle_request( + logical_region_id, + RegionRequest::Create(region_create_request), + ) + .await + .unwrap(); + + // Writing microsecond rows works. + let affected_rows = engine + .handle_request( + logical_region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: test_util::row_schema_with_tags_and_ts_datatype( + &["job"], + ColumnDataType::TimestampMicrosecond, + ), + rows: test_util::build_rows_with_ts_datatype( + 1, + 2, + ColumnDataType::TimestampMicrosecond, + ), + }, + hint: None, + partition_expr_version: None, + skip_wal: false, + }), + ) + .await + .unwrap(); + assert_eq!(affected_rows.affected_rows, 2); + + // The scan returns timestamps in the physical region's unit. + let stream = engine + .scan_to_stream(logical_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 timestamps = batch + .column(batch.schema().index_of(greptime_timestamp()).unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(batch.schema().index_of(greptime_value()).unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend( + timestamps + .values() + .iter() + .copied() + .zip(values.values().iter().copied()), + ); + } + rows.sort_unstable_by_key(|(timestamp, _)| *timestamp); + assert_eq!(rows, vec![(0, 0.0), (1, 1.0)]); + } + async fn create_logical_region_with_tags( env: &TestEnv, physical_region_id: RegionId, diff --git a/src/metric-engine/src/test_util.rs b/src/metric-engine/src/test_util.rs index 1f177370988..cf64d968182 100644 --- a/src/metric-engine/src/test_util.rs +++ b/src/metric-engine/src/test_util.rs @@ -121,6 +121,17 @@ impl TestEnv { self.metric.clone() } + /// Restarts Mito without closing regions and recreates the metric metadata cache. + /// Callers stop the old metric engine's metadata flush task before restarting. + pub async fn reopen_engine(&mut self, config: EngineConfig) { + let mito_config = self.mito.mito_config().clone(); + self.mito = self + .mito_env + .reopen_engine(self.mito.clone(), mito_config) + .await; + self.metric = MetricEngine::try_new(self.mito.clone(), config).unwrap(); + } + /// Creates a new follower engine with the same config as the leader engine. pub async fn create_follower_engine(&mut self) -> (MitoEngine, MetricEngine) { let mito = self @@ -158,6 +169,24 @@ impl TestEnv { physical_region_id: RegionId, table_dir: &str, options: Vec<(String, String)>, + ) { + self.create_physical_region_with_ts_type( + physical_region_id, + table_dir, + options, + ConcreteDataType::timestamp_millisecond_datatype(), + ) + .await + } + + /// Create regions in [MetricEngine] with specific `physical_region_id` + /// and time index column type. + pub async fn create_physical_region_with_ts_type( + &self, + physical_region_id: RegionId, + table_dir: &str, + options: Vec<(String, String)>, + ts_type: ConcreteDataType, ) { let region_create_request = RegionCreateRequest { engine: METRIC_ENGINE_NAME.to_string(), @@ -165,11 +194,7 @@ impl TestEnv { ColumnMetadata { column_id: 0, semantic_type: SemanticType::Timestamp, - column_schema: ColumnSchema::new( - greptime_timestamp(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), + column_schema: ColumnSchema::new(greptime_timestamp(), ts_type, false), }, ColumnMetadata { column_id: 1, @@ -330,16 +355,28 @@ pub fn create_logical_region_request( tags: &[&str], physical_region_id: RegionId, table_dir: &str, +) -> RegionCreateRequest { + create_logical_region_request_with_ts_type( + tags, + physical_region_id, + table_dir, + ConcreteDataType::timestamp_millisecond_datatype(), + ) +} + +/// Generate a [RegionCreateRequest] for logical region with the given time +/// index column type. +pub fn create_logical_region_request_with_ts_type( + tags: &[&str], + physical_region_id: RegionId, + table_dir: &str, + ts_type: ConcreteDataType, ) -> RegionCreateRequest { let mut column_metadatas = vec![ ColumnMetadata { column_id: 0, semantic_type: SemanticType::Timestamp, - column_schema: ColumnSchema::new( - greptime_timestamp(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), + column_schema: ColumnSchema::new(greptime_timestamp(), ts_type, false), }, ColumnMetadata { column_id: 1, @@ -407,10 +444,20 @@ pub fn alter_logical_region_request(tags: &[&str]) -> RegionAlterRequest { /// /// The result will also contains default timestamp and value column at beginning. pub fn row_schema_with_tags(tags: &[&str]) -> Vec { + row_schema_with_tags_and_ts_datatype(tags, ColumnDataType::TimestampMillisecond) +} + +/// Generate a row schema with given tag columns and time index datatype. +/// +/// The result will also contains default timestamp and value column at beginning. +pub fn row_schema_with_tags_and_ts_datatype( + tags: &[&str], + ts_datatype: ColumnDataType, +) -> Vec { let mut schema = vec![ PbColumnSchema { column_name: greptime_timestamp().to_string(), - datatype: ColumnDataType::TimestampMillisecond as i32, + datatype: ts_datatype as i32, semantic_type: SemanticType::Timestamp as _, datatype_extension: None, options: None, @@ -440,11 +487,33 @@ pub fn row_schema_with_tags(tags: &[&str]) -> Vec { /// The schema is generated by [row_schema_with_tags]. `num_tags` doesn't need to be precise, /// it's used to determine the column id for new columns. pub fn build_rows(num_tags: usize, num_rows: usize) -> Vec { + build_rows_with_ts_datatype(num_tags, num_rows, ColumnDataType::TimestampMillisecond) +} + +/// Build [Row]s whose time index values use the given timestamp datatype. +/// +/// The schema is generated by [row_schema_with_tags_and_ts_datatype]. +/// `num_tags` doesn't need to be precise, it's used to determine the column id +/// for new columns. +pub fn build_rows_with_ts_datatype( + num_tags: usize, + num_rows: usize, + ts_datatype: ColumnDataType, +) -> Vec { + let ts_value = |i: usize| -> ValueData { + match ts_datatype { + ColumnDataType::TimestampSecond => ValueData::TimestampSecondValue(i as _), + ColumnDataType::TimestampMillisecond => ValueData::TimestampMillisecondValue(i as _), + ColumnDataType::TimestampMicrosecond => ValueData::TimestampMicrosecondValue(i as _), + ColumnDataType::TimestampNanosecond => ValueData::TimestampNanosecondValue(i as _), + _ => unreachable!("not a timestamp datatype"), + } + }; let mut rows = vec![]; for i in 0..num_rows { let mut values = vec![ Value { - value_data: Some(ValueData::TimestampMillisecondValue(i as _)), + value_data: Some(ts_value(i)), }, Value { value_data: Some(ValueData::F64Value(i as f64)), diff --git a/src/mito-codec/src/key_values.rs b/src/mito-codec/src/key_values.rs index d66110bacfa..aba3cff988c 100644 --- a/src/mito-codec/src/key_values.rs +++ b/src/mito-codec/src/key_values.rs @@ -98,61 +98,6 @@ impl KeyValues { } } -/// Key value view of a mutation. -#[derive(Debug)] -pub struct KeyValuesRef<'a> { - /// Mutation to read. - /// - /// This mutation must be a valid mutation and rows in the mutation - /// must not be `None`. - mutation: &'a Mutation, - /// Key value read helper. - helper: SparseReadRowHelper, - /// Primary key encoding hint. - primary_key_encoding: PrimaryKeyEncoding, -} - -impl<'a> KeyValuesRef<'a> { - /// Creates [crate::memtable::KeyValues] from specific `mutation`. - /// - /// Returns `None` if `rows` of the `mutation` is `None`. - pub fn new(metadata: &RegionMetadata, mutation: &'a Mutation) -> Option> { - let rows = mutation.rows.as_ref()?; - let primary_key_encoding = - infer_primary_key_encoding_from_hint(mutation.write_hint.as_ref()); - let helper = SparseReadRowHelper::new(metadata, rows, primary_key_encoding); - - Some(KeyValuesRef { - mutation, - helper, - primary_key_encoding, - }) - } - - /// Returns a key value iterator. - pub fn iter(&self) -> impl Iterator> { - let rows = self.mutation.rows.as_ref().unwrap(); - let schema = &rows.schema; - rows.rows.iter().enumerate().map(|(idx, row)| { - KeyValue { - row, - schema, - helper: &self.helper, - sequence: self.mutation.sequence + idx as u64, // Calculate sequence for each row. - // Safety: This is a valid mutation. - op_type: OpType::try_from(self.mutation.op_type).unwrap(), - primary_key_encoding: self.primary_key_encoding, - } - }) - } - - /// Returns number of rows. - pub fn num_rows(&self) -> usize { - // Safety: rows is not None. - self.mutation.rows.as_ref().unwrap().rows.len() - } -} - /// Key value view of a row. /// /// A key value view divides primary key columns and field (value) columns. diff --git a/src/mito-codec/src/row_converter.rs b/src/mito-codec/src/row_converter.rs index fae2997182c..60bbf1935e0 100644 --- a/src/mito-codec/src/row_converter.rs +++ b/src/mito-codec/src/row_converter.rs @@ -134,6 +134,11 @@ pub trait PrimaryKeyCodec: Send + Sync + Debug { /// Returns the encoding type of the primary key. fn encoding(&self) -> PrimaryKeyEncoding; + /// Returns the dense codec for schema-aware positional access, if applicable. + fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> { + None + } + /// Decodes the primary key from the given bytes. /// /// Returns a [`CompositeValues`] that follows the primary key ordering. @@ -166,3 +171,28 @@ pub fn build_primary_key_codec_with_fields( } } } + +/// Finds the checked boundary of an Option, shared by Dense and Sparse. +/// This validates framing, not UTF-8; consumers decoding strings validate UTF-8. +pub(crate) fn encoded_string_len(bytes: &[u8]) -> memcomparable::Result { + match bytes.first().copied().ok_or(memcomparable::Error::Eof)? { + 0 => return Ok(1), + 1 => {} + marker => return Err(memcomparable::Error::InvalidTagEncoding(marker as usize)), + } + match bytes.get(1).copied().ok_or(memcomparable::Error::Eof)? { + 0 => return Ok(2), + 1 => {} + marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)), + } + let mut end = 2; + loop { + let chunk = bytes.get(end..end + 9).ok_or(memcomparable::Error::Eof)?; + end += 9; + match chunk[8] { + 1..=8 => return Ok(end), + 9 => {} + marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)), + } + } +} diff --git a/src/mito-codec/src/row_converter/dense.rs b/src/mito-codec/src/row_converter/dense.rs index 56289a08b11..ad85f4ac8f1 100644 --- a/src/mito-codec/src/row_converter/dense.rs +++ b/src/mito-codec/src/row_converter/dense.rs @@ -27,7 +27,7 @@ use datatypes::value::ValueRef; use memcomparable::{Deserializer, Serializer}; use paste::paste; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::ColumnId; @@ -38,7 +38,7 @@ use crate::error::{ use crate::key_values::KeyValue; use crate::primary_key_filter::DensePrimaryKeyFilter; use crate::row_converter::{ - CompositeValues, PrimaryKeyCodec, PrimaryKeyCodecExt, PrimaryKeyFilter, + CompositeValues, PrimaryKeyCodec, PrimaryKeyCodecExt, PrimaryKeyFilter, encoded_string_len, }; /// Field to serialize and deserialize value in memcomparable format. @@ -290,84 +290,112 @@ impl SortField { bytes: &[u8], deserializer: &mut Deserializer<&[u8]>, ) -> Result { - let len = self.encoded_field_len(&bytes[deserializer.position()..])?; + let pos = deserializer.position(); + let remaining = bytes + .get(pos..) + .ok_or(memcomparable::Error::Eof) + .context(error::DeserializeFieldSnafu)?; + let len = Self::encoded_len(self.encode_data_type(), remaining)?; + if len > 1 && self.encode_data_type().is_boolean() && remaining[1] > 1 { + return Err(memcomparable::Error::InvalidBoolEncoding(remaining[1])) + .context(error::DeserializeFieldSnafu); + } deserializer.advance(len); Ok(len) } /// Checks field boundaries before any unchecked reads in memcomparable. - fn encoded_field_len(&self, bytes: &[u8]) -> Result { - let invalid = || { - error::InvalidDensePrimaryKeySnafu { - reason: "truncated field or invalid encoding", + fn encoded_len(data_type: &ConcreteDataType, bytes: &[u8]) -> Result { + let marker = bytes + .first() + .copied() + .ok_or(memcomparable::Error::Eof) + .context(error::DeserializeFieldSnafu)?; + match marker { + 0 => return Ok(1), + 1 => {} + value => { + return Err(memcomparable::Error::InvalidTagEncoding(value as usize)) + .context(error::DeserializeFieldSnafu); } - .build() - }; - match bytes.first() { - Some(0) => return Ok(1), - Some(1) => {} - _ => return Err(invalid()), } - let len = match self.encode_data_type() { - ConcreteDataType::Boolean(_) - | ConcreteDataType::Int8(_) - | ConcreteDataType::UInt8(_) => 2, + let to_skip = match data_type { + ConcreteDataType::Boolean(_) => 2, + ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => 2, ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => 3, - ConcreteDataType::Int32(_) - | ConcreteDataType::UInt32(_) - | ConcreteDataType::Float32(_) - | ConcreteDataType::Date(_) => 5, - ConcreteDataType::Int64(_) - | ConcreteDataType::UInt64(_) - | ConcreteDataType::Float64(_) - | ConcreteDataType::Timestamp(_) => 9, - ConcreteDataType::Time(_) | ConcreteDataType::Duration(_) => 10, + ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => 5, + ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => 9, + ConcreteDataType::Float32(_) => 5, + ConcreteDataType::Float64(_) => 9, + ConcreteDataType::Binary(_) + | ConcreteDataType::Json(_) + | ConcreteDataType::Vector(_) => { + // Binary is encoded as a sequence of bytes, not chunked strings. + return encoded_binary_len(bytes).context(error::DeserializeFieldSnafu); + } + ConcreteDataType::String(_) => { + return encoded_string_len(bytes).context(error::DeserializeFieldSnafu); + } + ConcreteDataType::Date(_) => 5, + ConcreteDataType::Timestamp(_) => 9, // We treat timestamp as Option + ConcreteDataType::Time(_) => 10, // i64 and 1 byte time unit + ConcreteDataType::Duration(_) => 10, ConcreteDataType::Interval(IntervalType::YearMonth(_)) => 5, ConcreteDataType::Interval(IntervalType::DayTime(_)) => 9, ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => 17, ConcreteDataType::Decimal128(_) => 19, - ConcreteDataType::Binary(_) - | ConcreteDataType::Json(_) - | ConcreteDataType::Vector(_) => { - // A serde sequence: each byte is preceded by 1, terminated by 0. - let mut pos = 1; - while bytes.get(pos) == Some(&1) { - pos += 2; - } - if bytes.get(pos) != Some(&0) { - return Err(invalid()); - } - pos + 1 - } - ConcreteDataType::String(_) => { - match bytes.get(1) { - Some(0) => return Ok(2), - Some(1) => {} - _ => return Err(invalid()), - } - let mut pos = 2; - loop { - match bytes.get(pos + 8) { - Some(1..=8) => break pos + 9, - Some(9) => pos += 9, - _ => return Err(invalid()), - } - } - } - data_type => { - return error::NotSupportedFieldSnafu { + ConcreteDataType::Null(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Dictionary(_) => { + return NotSupportedFieldSnafu { data_type: data_type.clone(), } .fail(); } }; - snafu::ensure!( - bytes.len() >= len, - error::InvalidDensePrimaryKeySnafu { - reason: "truncated field", - } - ); - Ok(len) + if bytes.len() < to_skip { + return Err(memcomparable::Error::Eof).context(error::DeserializeFieldSnafu); + } + Ok(to_skip) + } + + /// Decodes primitive fields straight into a reference value for column + /// builders, avoiding owned Value construction and conversion. The caller has + /// already checked the encoded boundary. + fn deserialize_primitive_ref( + data_type: &ConcreteDataType, + deserializer: &mut Deserializer, + ) -> Option>> { + macro_rules! decode_primitive { + ($($variant:ident, $native:ty),* $(,)?) => { + match data_type { + $(ConcreteDataType::$variant(_) => Some( + Option::<$native>::deserialize(deserializer) + .map(ValueRef::from) + .context(error::DeserializeFieldSnafu) + ),)* + _ => None, + } + }; + } + decode_primitive!( + Boolean, bool, Int8, i8, Int16, i16, Int32, i32, Int64, i64, UInt8, u8, UInt16, u16, + UInt32, u32, UInt64, u64, Float32, f32, Float64, f64, + ) + } +} + +/// The Option marker has already been checked by SortField. +fn encoded_binary_len(bytes: &[u8]) -> memcomparable::Result { + let mut current = 1; + loop { + match bytes.get(current).copied() { + Some(0) => return Ok(current + 1), + Some(1) if bytes.get(current + 1).is_some() => current += 2, + Some(1) | None => return Err(memcomparable::Error::Eof), + Some(marker) => return Err(memcomparable::Error::InvalidSeqEncoding(marker)), + } } } @@ -444,7 +472,16 @@ impl DensePrimaryKeyCodec { return Ok(index); } let start = deserializer.position(); - let len = field.encoded_field_len(&bytes[start..])?; + // Preserve the prefix error contract used by primary-key range mapping. + let len = SortField::encoded_len(field.encode_data_type(), &bytes[start..]).map_err( + |source| match source { + error::Error::DeserializeField { .. } => error::InvalidDensePrimaryKeySnafu { + reason: "truncated field or invalid encoding", + } + .build(), + source => source, + }, + )?; // Production range mapping only needs boundaries, not allocated field values. #[cfg(any(debug_assertions, test))] field.deserialize(&mut Deserializer::new(&bytes[start..start + len]))?; @@ -470,6 +507,73 @@ impl DensePrimaryKeyCodec { Ok(values) } + /// Iterates over all source PK fields in schema order, checking each field's + /// encoded boundary before deserializing it. Stops after the first error. + /// + /// Unlike positional access, this requires no offsets or value cache. Callers + /// can consume each value immediately instead of retaining the entire key. + pub fn decode_dense_iter<'a>( + &'a self, + bytes: &'a [u8], + ) -> impl Iterator> + 'a { + self.ordered_primary_key_columns.iter().scan( + Some(Deserializer::new(bytes)), + move |state, (id, field)| { + let deserializer = state.as_mut()?; + let remaining = &bytes[deserializer.position()..]; + let decoded = SortField::encoded_len(field.encode_data_type(), remaining) + .and_then(|_| field.deserialize(deserializer).map(|value| (*id, value))); + if decoded.is_err() { + *state = None; + } + Some(decoded) + }, + ) + } + + /// Returns the column ids and sort fields in encoded primary-key order. + pub fn fields(&self) -> &[(ColumnId, SortField)] { + &self.ordered_primary_key_columns + } + + /// Decodes all fields in source order into a consumer, checking boundaries. + /// Strings borrow the reusable buffer for the duration of each callback, so + /// materializing them into column builders needs no per-value owned string. + pub fn decode_dense_with( + &self, + bytes: &[u8], + value_buf: &mut Vec, + mut consume: impl FnMut(usize, ValueRef<'_>), + ) -> Result<()> { + let mut deserializer = Deserializer::new(bytes); + for (pos, (_, field)) in self.ordered_primary_key_columns.iter().enumerate() { + let data_type = field.encode_data_type(); + let remaining = &bytes[deserializer.position()..]; + SortField::encoded_len(data_type, remaining)?; + if data_type.is_string() && remaining[0] != 0 { + deserializer.advance(1); + deserializer + .read_bytes_into(value_buf) + .context(error::DeserializeFieldSnafu)?; + let value = std::str::from_utf8(value_buf).map_err(|err| { + error::InvalidDensePrimaryKeySnafu { + reason: format!("string is not valid UTF-8: {err}"), + } + .build() + })?; + consume(pos, ValueRef::String(value)); + } else if let Some(value) = + SortField::deserialize_primitive_ref(data_type, &mut deserializer) + { + consume(pos, value?); + } else { + let value = field.deserialize(&mut deserializer)?; + consume(pos, value.as_value_ref()); + } + } + Ok(()) + } + /// Returns the field at `pos`. /// /// # Panics @@ -488,61 +592,61 @@ impl DensePrimaryKeyCodec { offsets_buf: &mut Vec, deserializer: &mut Deserializer<&[u8]>, ) -> Result { - if pos < offsets_buf.len() { - // We computed the offset before. - let offset = offsets_buf[pos]; - deserializer.advance(offset); - return Ok(offset); - } - + ensure!( + pos < self.num_fields(), + error::InvalidDensePrimaryKeySnafu { + reason: format!( + "field position {pos} exceeds field count {}", + self.num_fields() + ), + } + ); if offsets_buf.is_empty() { - let mut offset = 0; - // Skip values before `pos`. - for i in 0..pos { - // Offset to skip before reading value i. - offsets_buf.push(offset); - let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?; - offset += skip; - } - // Offset to skip before reading this value. - offsets_buf.push(offset); - Ok(offset) - } else { - // Offsets are not enough. - let value_start = offsets_buf.len() - 1; - // Advances to decode value at `value_start`. - let mut offset = offsets_buf[value_start]; - deserializer.advance(offset); - for i in value_start..pos { - // Skip value i. - let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?; - // Offset for the value at i + 1. - offset += skip; - offsets_buf.push(offset); - } - Ok(offset) + // The schema bounds this cache. Reserve once even when the first + // requested field is near the end of a wide key. + offsets_buf.reserve_exact(self.num_fields()); + offsets_buf.push(0); } + // Start at the requested field if cached, otherwise resume discovery + // from the furthest known boundary. + let value_start = pos.min(offsets_buf.len() - 1); + let mut offset = offsets_buf[value_start]; + ensure!( + offset <= bytes.len(), + error::InvalidDensePrimaryKeySnafu { + reason: "cached offset exceeds key length", + } + ); + deserializer.advance(offset); + for i in value_start..pos { + offset += self.field_at(i).skip_deserialize(bytes, deserializer)?; + offsets_buf.push(offset); + } + Ok(offset) } /// Decode value at `pos` in `bytes`. /// - /// The i-th element in offsets buffer is how many bytes to skip in order to read value at `pos`. + /// The i-th element in the offsets buffer is the start of field i. The buffer + /// must only be reused for the same encoded key and codec. + #[inline] pub fn decode_value_at( &self, bytes: &[u8], pos: usize, offsets_buf: &mut Vec, ) -> Result { - let mut deserializer = Deserializer::new(bytes); - self.advance_to_value_at(bytes, pos, offsets_buf, &mut deserializer)?; - - self.field_at(pos).deserialize(&mut deserializer) + let encoded = self.encoded_value_at(bytes, pos, offsets_buf)?; + self.field_at(pos) + .deserialize(&mut Deserializer::new(encoded)) } /// Returns the encoded bytes at `pos` in `bytes`. /// - /// The i-th element in offsets buffer is how many bytes to skip in order to read value at - /// `pos`. + /// The i-th element in the offsets buffer is the start of field i. The buffer + /// must only be reused for the same encoded key and codec. Validates the + /// framing up to this field, without validating or decoding later fields. + #[inline] pub fn encoded_value_at<'a>( &self, bytes: &'a [u8], @@ -555,6 +659,11 @@ impl DensePrimaryKeyCodec { let len = self .field_at(pos) .skip_deserialize(bytes, &mut deserializer)?; + // We already found the next field's start while validating this one. + // Reuse it instead of scanning this field again on the next lookup. + if offsets_buf.len() == pos + 1 && pos + 1 < self.num_fields() { + offsets_buf.push(offset + len); + } Ok(&bytes[offset..offset + len]) } @@ -600,6 +709,10 @@ impl PrimaryKeyCodec for DensePrimaryKeyCodec { PrimaryKeyEncoding::Dense } + fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> { + Some(self) + } + fn primary_key_filter( &self, metadata: &RegionMetadataRef, @@ -663,6 +776,78 @@ mod tests { } let decoded = encoder.decode(&result).unwrap().into_dense(); assert_eq!(decoded, row); + let mut value_buf = Vec::new(); + let mut borrowed = Vec::new(); + encoder + .decode_dense_with(&result, &mut value_buf, |_, value| { + borrowed.push(Value::from(value)) + }) + .unwrap(); + assert_eq!(borrowed, row); + let sequential = encoder + .decode_dense_iter(&result) + .collect::>>() + .unwrap(); + assert_eq!( + sequential + .iter() + .map(|(_, value)| value) + .collect::>(), + row.iter().collect::>() + ); + for end in 0..result.len() { + assert!( + encoder + .decode_dense_with(&result[..end], &mut value_buf, |_, _| {}) + .is_err() + ); + assert!( + encoder + .decode_dense_iter(&result[..end]) + .collect::>>() + .is_err(), + "truncated at {end}" + ); + } + // Every supported type must have the same boundaries in positional and + // full decoding. Exercise every truncation both at the requested field + // and while skipping it to reach a later field. + let mut offsets = Vec::new(); + let mut end = 0; + for (pos, value) in row.iter().enumerate() { + let encoded = encoder + .encoded_value_at(&result, pos, &mut offsets) + .unwrap(); + assert_eq!( + encoder + .field_at(pos) + .deserialize(&mut Deserializer::new(encoded)) + .unwrap(), + *value + ); + let start = end; + end += encoded.len(); + for len in start..end { + assert!( + encoder + .encoded_value_at(&result[..len], pos, &mut Vec::new()) + .is_err() + ); + assert!( + encoder + .decode_value_at(&result[..len], pos, &mut Vec::new()) + .is_err() + ); + if pos + 1 < row.len() { + assert!( + encoder + .encoded_value_at(&result[..len], pos + 1, &mut Vec::new()) + .is_err() + ); + } + } + } + assert_eq!(end, result.len()); let mut decoded = Vec::new(); let mut offsets = Vec::new(); // Iter two times to test offsets buffer. @@ -871,6 +1056,96 @@ mod tests { } } + #[test] + fn test_positional_decode_invalid_encoding() { + for (ty, bytes) in [ + (ConcreteDataType::int64_datatype(), vec![2]), + (ConcreteDataType::boolean_datatype(), vec![1, 2]), + (ConcreteDataType::binary_datatype(), vec![1, 2]), + (ConcreteDataType::binary_datatype(), vec![1, 1, 42, 2]), + (ConcreteDataType::string_datatype(), vec![1, 2]), + ( + ConcreteDataType::string_datatype(), + vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + ), + ( + ConcreteDataType::string_datatype(), + vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 10], + ), + ] { + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (0, SortField::new(ty)), + (1, SortField::new(ConcreteDataType::int64_datatype())), + ]); + for pos in [0, 1, 2, usize::MAX] { + assert!( + codec + .encoded_value_at(&bytes, pos, &mut Vec::new()) + .is_err() + ); + assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err()); + } + assert!( + codec + .decode_dense_iter(&bytes) + .collect::>>() + .is_err() + ); + assert!( + codec + .decode_dense_with(&bytes, &mut Vec::new(), |_, _| {}) + .is_err() + ); + } + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (0, SortField::new(ConcreteDataType::string_datatype())), + (1, SortField::new(ConcreteDataType::int64_datatype())), + ]); + let bytes = codec + .encode([ValueRef::String("abcdefghijk"), ValueRef::Int64(42)].into_iter()) + .unwrap(); + for pos in [2, usize::MAX] { + assert!( + codec + .encoded_value_at(&bytes, pos, &mut Vec::new()) + .is_err() + ); + assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err()); + } + for offset in [bytes.len(), bytes.len() + 1, usize::MAX] { + for pos in [0, 1] { + assert!( + codec + .encoded_value_at(&bytes, pos, &mut vec![offset]) + .is_err() + ); + assert!( + codec + .decode_value_at(&bytes, pos, &mut vec![offset]) + .is_err() + ); + } + } + let mut invalid_utf8 = bytes; + invalid_utf8[2] = 0xff; + assert!( + codec + .decode_dense_with(&invalid_utf8, &mut Vec::new(), |_, _| {}) + .is_err() + ); + assert!( + codec + .decode_dense_iter(&invalid_utf8) + .collect::>>() + .is_err() + ); + assert!( + codec + .decode_value_at(&invalid_utf8, 0, &mut Vec::new()) + .is_err() + ); + } + #[test] fn test_memcmp_dictionary() { // Test Dictionary diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index 1be2fbacb2b..5497a7dc034 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/src/mito-codec/src/row_converter/sparse.rs @@ -581,40 +581,6 @@ impl PrimaryKeyCodec for SparsePrimaryKeyCodec { } } -/// Field with column id. -pub struct FieldWithId { - pub field: SortField, - pub column_id: ColumnId, -} - -/// A special encoder for memtable. -pub struct SparseEncoder { - fields: Vec, -} - -impl SparseEncoder { - pub fn new(fields: Vec) -> Self { - Self { fields } - } - - pub fn encode_to_vec<'a, I>(&self, row: I, buffer: &mut Vec) -> Result<()> - where - I: Iterator>, - { - let mut serializer = Serializer::new(buffer); - for (value, field) in row.zip(self.fields.iter()) { - if !value.is_null() { - field - .column_id - .serialize(&mut serializer) - .context(SerializeFieldSnafu)?; - field.field.serialize(&mut serializer, &value)?; - } - } - Ok(()) - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs index 1e57c13cc7f..fec89b74e0d 100644 --- a/src/mito-codec/src/row_converter/sparse/checked.rs +++ b/src/mito-codec/src/row_converter/sparse/checked.rs @@ -17,6 +17,7 @@ use snafu::{OptionExt, ensure}; use store_api::storage::ColumnId; use crate::error::{InvalidSparsePrimaryKeySnafu, Result}; +use crate::row_converter::encoded_string_len; use crate::row_converter::sparse::{ COLUMN_ID_ENCODE_SIZE, RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparseOffsetsCache, TABLE_ID_VALUE_OFFSET, TAGS_START_OFFSET, TSID_VALUE_OFFSET, @@ -115,7 +116,7 @@ impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> { // encoded_value checked the markers and every chunk length. Reserve once // before unchunking so the buffer does not grow for each 8-byte chunk. buffer.reserve(encoded.len() - 2); - for chunk in encoded[2..].chunks_exact(9) { + for chunk in encoded[2..].as_chunks::<9>().0 { let len = usize::from(chunk[8]).min(8); buffer.extend_from_slice(&chunk[..len]); } @@ -141,45 +142,13 @@ impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> { /// Finds the end of an Option without allocating or reading its payload. /// Unlike memcomparable's skip_bytes, all advances are checked for truncated input. fn encoded_label(bytes: &[u8]) -> Result<&[u8]> { - match bytes.first() { - Some(0) => return Ok(&bytes[..1]), - Some(1) => {} - _ => { - return InvalidSparsePrimaryKeySnafu { - reason: "invalid label null marker", - } - .fail(); + let len = encoded_string_len(bytes).map_err(|error| { + InvalidSparsePrimaryKeySnafu { + reason: error.to_string(), } - } - match bytes.get(1) { - Some(0) => return Ok(&bytes[..2]), - Some(1) => {} - _ => { - return InvalidSparsePrimaryKeySnafu { - reason: "invalid label bytes marker", - } - .fail(); - } - } - let mut end = 2; - loop { - let chunk = bytes - .get(end..end + 9) - .context(InvalidSparsePrimaryKeySnafu { - reason: "truncated label chunk", - })?; - end += 9; - match chunk[8] { - 1..=8 => return Ok(&bytes[..end]), - 9 => {} - _ => { - return InvalidSparsePrimaryKeySnafu { - reason: "invalid label chunk length", - } - .fail(); - } - } - } + .build() + })?; + Ok(&bytes[..len]) } #[cfg(test)] diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 583ec31bd4a..b8fd2f8892c 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, bucket planning, catalogs, worker-owned maintenance and purge tasks, and index lifecycle | +| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning, catalogs, worker-owned reconciliation (`maintenance.rs`), shared approximate usage (`task.rs`), catalogs, and index lifecycle | | `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,10 @@ 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 disk limit**: tasks estimate usage from entry sizes in open-region snapshots + and defer new builds when the shared estimate is full. Cleanup still prunes obsolete + range entries and expires series indexes to reclaim published usage. Concurrent builds may + overshoot; closed-region files are not counted. Series handles and the SST purger own file deletion. - **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`. diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 2be72af57cf..39dd96f687b 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -10,7 +10,7 @@ test = ["common-test-util", "rstest", "rstest_reuse", "rskafka"] testing = ["test"] test-shared-fs-region-migration = [] enterprise = [] -vector_index = ["dep:roaring", "index/vector_index"] +hdfs-object-store = ["object-store/hdfs-object-store"] [lints] workspace = true @@ -20,7 +20,6 @@ api.workspace = true aquamarine.workspace = true arrow-schema.workspace = true async-channel = "1.9" -common-stat.workspace = true async-stream.workspace = true async-trait.workspace = true base64.workspace = true @@ -28,58 +27,54 @@ bytemuck.workspace = true bytes.workspace = true chrono.workspace = true common-base.workspace = true -common-config.workspace = true common-datasource.workspace = true -common-decimal.workspace = true common-error.workspace = true -common-grpc.workspace = true common-function.workspace = true +common-grpc.workspace = true common-macro.workspace = true common-memory-manager.workspace = true common-meta.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true +common-stat.workspace = true common-telemetry.workspace = true common-test-util = { workspace = true, optional = true } common-time.workspace = true common-wal.workspace = true crc32fast = "1" crossbeam-utils.workspace = true +dashmap.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true datatypes.workspace = true -dashmap.workspace = true derive_more.workspace = true dotenv.workspace = true either.workspace = true futures.workspace = true +greptime-proto.workspace = true humantime.workspace = true humantime-serde.workspace = true index.workspace = true itertools.workspace = true -greptime-proto.workspace = true lazy_static = "1.4" log-store = { workspace = true } mito-codec.workspace = true moka = { workspace = true, features = ["sync", "future"] } object-store = { workspace = true, features = ["testing"] } parquet = { workspace = true, features = ["async", "variant_experimental"] } -paste.workspace = true +partition.workspace = true pin-project.workspace = true prometheus.workspace = true prost.workspace = true -partition.workspace = true puffin.workspace = true rand.workspace = true rayon = "1.10" -roaring = { version = "0.10", optional = true } regex.workspace = true rskafka = { workspace = true, optional = true } rstest = { workspace = true, optional = true } rstest_reuse = { workspace = true, optional = true } -scopeguard = "1.2" serde.workspace = true serde_json.workspace = true serde_with.workspace = true @@ -102,13 +97,12 @@ common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true criterion = { workspace = true, features = ["async", "async_tokio"] } dotenv.workspace = true -log-store.workspace = true +log-store = { workspace = true, features = ["testing"] } mito-codec = { workspace = true, features = ["testing"] } object-store = { workspace = true, features = ["services-memory"] } rskafka.workspace = true rstest.workspace = true rstest_reuse.workspace = true -session.workspace = true toml.workspace = true [[bench]] @@ -159,6 +153,3 @@ required-features = ["testing"] name = "bench_dense_index_update" harness = false required-features = ["testing"] - -[package.metadata.cargo-udeps.ignore] -normal = ["aquamarine"] diff --git a/src/mito2/benches/bench_dense_index_update.rs b/src/mito2/benches/bench_dense_index_update.rs index f81d30eb4a6..6f5ccfc080b 100644 --- a/src/mito2/benches/bench_dense_index_update.rs +++ b/src/mito2/benches/bench_dense_index_update.rs @@ -30,8 +30,8 @@ use api::v1::SemanticType; use common_test_util::temp_dir::create_temp_dir; use criterion::{Criterion, Throughput, criterion_group, criterion_main}; use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, StringDictionaryBuilder, TimestampMillisecondArray, - UInt8Array, UInt64Array, + ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, StringDictionaryBuilder, + TimestampMillisecondArray, UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::UInt32Type; use datatypes::arrow::record_batch::RecordBatch; @@ -39,6 +39,8 @@ use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; use datatypes::value::Value; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; +use mito2::read::{Batch, BatchBuilder}; +use mito2::sst::index::Indexer; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use mito2::test_util::bench_util::{BloomFilterIndexer, InvertedIndexer}; @@ -284,55 +286,115 @@ fn bench_dense_index_update(c: &mut Criterion) { } for shape in shapes { let (metadata, inverted_columns, batch) = input(&shape); - let run = |iterations| { - runtime.block_on(async { - let mut elapsed = Duration::ZERO; - for _ in 0..iterations { - let file_id = FileId::random(); - let mut inverted = shape.inverted.then(|| { - InvertedIndexer::new( - file_id, - &metadata, - intermediate.clone(), - None, - NonZeroUsize::new(SEGMENT_ROWS).unwrap(), - inverted_columns.clone(), - ) - }); - let mut bloom = if shape.bloom { - BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None) - .unwrap() - } else { - None - }; - if allocations { - ALLOCATIONS.set(Some((0, 0))); + for encoded_only in [false, true] { + let name = if encoded_only { + format!("encoded/{}", shape.name) + } else { + shape.name.clone() + }; + let encoded_batches: Vec = if encoded_only { + let pk = batch + .column_by_name("__primary_key") + .unwrap() + .as_any() + .downcast_ref::>() + .unwrap(); + let values = pk.values().as_any().downcast_ref::().unwrap(); + (0..ROWS) + .step_by(shape.rows_per_key) + .map(|row| { + let mut builder = + BatchBuilder::new(values.value(pk.keys().value(row) as usize).to_vec()); + builder + .timestamps_array( + batch + .column_by_name("ts") + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + builder + .sequences_array( + batch + .column_by_name("__sequence") + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + builder + .op_types_array( + batch + .column_by_name("__op_type") + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + if shape.field { + builder + .push_field_array( + shape.tags, + batch + .column_by_name(&format!("col_{}", shape.tags)) + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + } + builder.build().unwrap() + }) + .collect() + } else { + Vec::new() + }; + let run = |iterations| { + runtime.block_on(async { + let mut elapsed = Duration::ZERO; + for _ in 0..iterations { + let file_id = FileId::random(); + let inverted = shape.inverted.then(|| { + InvertedIndexer::new( + file_id, + &metadata, + intermediate.clone(), + None, + NonZeroUsize::new(SEGMENT_ROWS).unwrap(), + inverted_columns.clone(), + ) + }); + let bloom = if shape.bloom { + BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None) + .unwrap() + } else { + None + }; + let mut indexer = Indexer::for_bench(&metadata, inverted, bloom); + // Fresh caches for each iteration; both creators share each Batch. + let mut encoded_batches = encoded_batches.clone(); + if allocations { + ALLOCATIONS.set(Some((0, 0))); + } + let start = Instant::now(); + if encoded_only { + for batch in &mut encoded_batches { + indexer.update(batch).await; + } + } else { + indexer.update_flat(&batch).await; + } + elapsed += start.elapsed(); + if allocations { + report_allocations(&name); + } + indexer.abort().await; } - let start = Instant::now(); - if let Some(indexer) = &mut inverted { - indexer.update_flat(&batch).await.unwrap(); - } - if let Some(indexer) = &mut bloom { - indexer.update_flat(&batch).await.unwrap(); - } - elapsed += start.elapsed(); - if allocations { - report_allocations(&shape.name); - } - if let Some(indexer) = &mut inverted { - indexer.abort().await.unwrap(); - } - if let Some(indexer) = &mut bloom { - indexer.abort().await.unwrap(); - } - } - elapsed - }) - }; - if allocations { - run(1); - } else { - group.bench_function(&shape.name, |b| b.iter_custom(&run)); + elapsed + }) + }; + if allocations { + run(1); + } else { + group.bench_function(&name, |b| b.iter_custom(&run)); + } } } group.finish(); diff --git a/src/mito2/benches/bench_pk_tag_column.rs b/src/mito2/benches/bench_pk_tag_column.rs index d3ded3d0b05..9e191a69511 100644 --- a/src/mito2/benches/bench_pk_tag_column.rs +++ b/src/mito2/benches/bench_pk_tag_column.rs @@ -24,6 +24,7 @@ use std::hint::black_box; use std::sync::Arc; +use std::time::Duration; use api::v1::SemanticType; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; @@ -35,9 +36,12 @@ use datatypes::arrow::datatypes::UInt32Type; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; -use mito_codec::row_converter::SparsePrimaryKeyCodec; +use datatypes::value::Value; +use mito_codec::row_converter::{ + DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField, SparsePrimaryKeyCodec, +}; use mito2::sst::parquet::flat_format::decode_primary_keys; -use mito2::test_util::bench_util::tag_filter_for_bench; +use mito2::test_util::bench_util::{pk_materializer_for_bench, tag_filter_for_bench}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; use store_api::storage::RegionId; @@ -237,5 +241,112 @@ fn bench_pk_tag_filters(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_pk_tag_column, bench_pk_tag_filters); +/// Measures encoded-only Dense inputs, including the cost of discovering offsets. +fn bench_dense_pk_tag_column(c: &mut Criterion) { + const TAGS: usize = 40; + let mut group = c.benchmark_group("dense_pk_tag_column"); + group.sample_size(30); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + for numeric in [false, true] { + let ty = if numeric { + ConcreteDataType::uint64_datatype() + } else { + ConcreteDataType::string_datatype() + }; + let codec = DensePrimaryKeyCodec::with_fields( + (0..TAGS) + .map(|id| (id as u32, SortField::new(ty.clone()))) + .collect(), + ); + let mut metadata = RegionMetadataBuilder::new(RegionId::new(1, 1)); + for id in 0..TAGS { + metadata.push_column_metadata(ColumnMetadata { + column_id: id as u32, + column_schema: ColumnSchema::new(format!("tag_{id}"), ty.clone(), true), + semantic_type: SemanticType::Tag, + }); + } + metadata + .push_column_metadata(ColumnMetadata { + column_id: TAGS as u32, + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + }) + .primary_key((0..TAGS as u32).collect()); + let metadata = Arc::new(metadata.build().unwrap()); + for rows_per_key in [1, 32] { + let mut keys = BinaryDictionaryBuilder::::new(); + for series in 0..ROWS / rows_per_key { + let values: Vec<_> = (0..TAGS) + .map(|id| { + if numeric { + Value::UInt64((series + id) as u64) + } else { + Value::from(format!( + "tag-{id:03}-value-{series:010}-abcdefghijklmnopqrstuvwxyz" + )) + } + }) + .collect(); + let pk = codec + .encode(values.iter().map(Value::as_value_ref)) + .unwrap(); + for _ in 0..rows_per_key { + keys.append(&pk).unwrap(); + } + } + let batch = RecordBatch::try_from_iter([ + ( + "ts", + Arc::new(TimestampMillisecondArray::from_iter_values(0..ROWS as i64)) + as ArrayRef, + ), + ("__primary_key", Arc::new(keys.finish()) as ArrayRef), + ( + "__sequence", + Arc::new(UInt64Array::from(vec![1; ROWS])) as ArrayRef, + ), + ( + "__op_type", + Arc::new(UInt8Array::from(vec![1; ROWS])) as ArrayRef, + ), + ]) + .unwrap(); + for (projection, start, count) in [ + ("first", 0, 1), + ("last", 39, 1), + ("4last", 36, 4), + ("all", 0, 40), + ] { + let kind = if numeric { "numeric" } else { "string" }; + let materialize = pk_materializer_for_bench( + metadata.clone(), + (start as u32..(start + count) as u32) + .chain([TAGS as u32]) + .collect(), + batch.schema(), + ); + assert_eq!(materialize(batch.clone()).unwrap().num_columns(), count + 4); + group.bench_function(format!("{kind}_{projection}_{rows_per_key}rpk"), |b| { + b.iter(|| { + black_box(materialize(black_box(batch.clone())).unwrap()); + }); + }); + } + } + } + group.finish(); +} + +criterion_group!( + benches, + bench_pk_tag_column, + bench_pk_tag_filters, + bench_dense_pk_tag_column +); criterion_main!(benches); diff --git a/src/mito2/benches/memtable_bench.rs b/src/mito2/benches/memtable_bench.rs index 80331941865..e4d71c9883c 100644 --- a/src/mito2/benches/memtable_bench.rs +++ b/src/mito2/benches/memtable_bench.rs @@ -268,7 +268,7 @@ fn filter_1_host(c: &mut Criterion) { ) .unwrap(); - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { let iter = range.build_record_batch_iter(None, None).unwrap(); for batch in iter { let _batch = batch.unwrap(); diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 337386aa2ff..3959e0c26a4 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -18,6 +18,7 @@ use std::time::{Duration, Instant}; use async_stream::try_stream; use common_base::readable_size::ReadableSize; +use common_runtime::runtime::RuntimeTrait; use common_telemetry::warn; use common_time::Timestamp; use futures::{Stream, TryStreamExt}; @@ -348,6 +349,7 @@ impl AccessLayer { write_opts: &WriteOptions, metrics: &mut Metrics, ) -> Result { + let op_type = request.op_type; let region_id = request.metadata.region_id; let region_metadata = request.metadata.clone(); let cache_manager = request.cache_manager.clone(); @@ -389,8 +391,6 @@ impl AccessLayer { inverted_index_config: request.inverted_index_config, fulltext_index_config: request.fulltext_index_config, bloom_filter_index_config: request.bloom_filter_index_config, - #[cfg(feature = "vector_index")] - vector_index_config: request.vector_index_config, }; // We disable write cache on file system but we still use atomic write. // TODO(yingwen): If we support other non-fs stores without the write cache, then @@ -426,6 +426,10 @@ impl AccessLayer { // Put parquet metadata to cache manager. if !sst_info.is_empty() && cache_manager.sst_meta_cache_enabled() { + let runtime = match op_type { + OperationType::Compact => common_runtime::compact_runtime(), + OperationType::Flush => common_runtime::global_runtime(), + }; for sst in &sst_info { if let Some(parquet_metadata) = &sst.file_metadata { let file_id = RegionFileId::new(region_id, sst.file_id); @@ -445,7 +449,7 @@ impl AccessLayer { // Compact cache preparation is best-effort. Run the entire operation in one // detached blocking task so it neither blocks an async worker nor delays the // SST write. - common_runtime::spawn_blocking_global(move || { + runtime.spawn_blocking(move || { match prepare_sst_meta_sync( &file_path, Arc::unwrap_or_clone(parquet_metadata), @@ -574,8 +578,6 @@ pub struct SstWriteRequest { pub metadata: RegionMetadataRef, pub source: FlatSource, pub cache_manager: CacheManagerRef, - #[allow(dead_code)] - pub storage: Option, /// Optional uniform row sequence for writes that do not preserve sequences. /// Compaction passes `None` to retain the reader's effective input sequences. pub max_sequence: Option, @@ -589,8 +591,6 @@ pub struct SstWriteRequest { pub inverted_index_config: InvertedIndexConfig, pub fulltext_index_config: FulltextIndexConfig, pub bloom_filter_index_config: BloomFilterConfig, - #[cfg(feature = "vector_index")] - pub vector_index_config: crate::config::VectorIndexConfig, } /// Cleaner to remove temp files on the atomic write dir. diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 0d43d31f434..cf0e8de3458 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -31,6 +31,8 @@ use std::sync::{Arc, RwLock, Weak}; use bytes::Bytes; use common_base::readable_size::ReadableSize; use common_datasource::compression::CompressionType; +use common_runtime::Runtime; +use common_runtime::runtime::RuntimeTrait; use common_telemetry::warn; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::record_batch::RecordBatch; @@ -56,8 +58,6 @@ pub use write_cache::{WriteCacheUploadStoreWrapper, WriteCacheUploadStoreWrapper use crate::cache::cache_size::parquet_meta_size; use crate::cache::file_cache::{FileType, IndexKey}; use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCacheRef}; -#[cfg(feature = "vector_index")] -use crate::cache::index::vector_index::{VectorIndexCache, VectorIndexCacheRef}; use crate::cache::write_cache::WriteCacheRef; use crate::error::{ CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, @@ -238,7 +238,7 @@ impl SstMetaPreparation { } impl CompactSstMeta { - async fn decode(&self) -> Result> { + async fn decode(&self, runtime: &Runtime) -> Result> { let mut decoded_guard = self.decoded.lock().await; if let Some(decoded) = decoded_guard.upgrade() { return Ok(decoded); @@ -248,33 +248,34 @@ impl CompactSstMeta { let decoded_size = self.decoded_size; let region_metadata = self.region_metadata.upgrade(); let page_index_policy = self.page_index_policy; - let decoded = common_runtime::spawn_blocking_global(move || { - let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context( - DecompressObjectSnafu { - compress_type: CompressionType::Zstd, + let decoded = runtime + .spawn_blocking(move || { + let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context( + DecompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: "cached SST metadata", + }, + )?; + let bytes = Bytes::from(bytes); + let mut reader = ParquetMetaDataReader::new() + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(page_index_policy); + reader.try_parse(&bytes).context(ReadParquetSnafu { path: "cached SST metadata", - }, - )?; - let bytes = Bytes::from(bytes); - let mut reader = ParquetMetaDataReader::new() - .with_column_index_policy(PageIndexPolicy::Skip) - .with_offset_index_policy(page_index_policy); - reader.try_parse(&bytes).context(ReadParquetSnafu { - path: "cached SST metadata", - })?; - let metadata = reader.finish().context(ReadParquetSnafu { - path: "cached SST metadata", - })?; - CachedSstMeta::try_new_with_page_index_policy( - "cached SST metadata", - metadata, - region_metadata, - page_index_policy, - ) - .map(Arc::new) - }) - .await - .context(JoinSnafu)??; + })?; + let metadata = reader.finish().context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + CachedSstMeta::try_new_with_page_index_policy( + "cached SST metadata", + metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)??; *decoded_guard = Arc::downgrade(&decoded); Ok(decoded) @@ -285,46 +286,50 @@ impl CompactSstMeta { } } -/// Decodes SST metadata without preparing a compact cache entry. +/// Decodes SST metadata on the given blocking runtime without preparing a compact cache entry. pub(crate) async fn decode_sst_meta( file_path: &str, parquet_metadata: ParquetMetaData, region_metadata: Option, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Result> { let file_path = file_path.to_string(); - common_runtime::spawn_blocking_global(move || { - let parquet_metadata = strip_column_indexes(parquet_metadata); - CachedSstMeta::try_new_with_page_index_policy( - &file_path, - parquet_metadata, - region_metadata, - page_index_policy, - ) - .map(Arc::new) - }) - .await - .context(JoinSnafu)? + runtime + .spawn_blocking(move || { + let parquet_metadata = strip_column_indexes(parquet_metadata); + CachedSstMeta::try_new_with_page_index_policy( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)? } -/// Decodes SST metadata and attempts to encode both cache representations on the blocking runtime. +/// Decodes SST metadata and attempts to encode both cache representations on the given blocking runtime. pub(crate) async fn prepare_sst_meta( file_path: &str, parquet_metadata: ParquetMetaData, region_metadata: Option, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Result { let file_path = file_path.to_string(); - common_runtime::spawn_blocking_global(move || { - prepare_sst_meta_sync( - &file_path, - parquet_metadata, - region_metadata, - page_index_policy, - ) - }) - .await - .context(JoinSnafu)? + runtime + .spawn_blocking(move || { + prepare_sst_meta_sync( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + }) + .await + .context(JoinSnafu)? } /// Synchronously prepares SST metadata. Callers must run this on a blocking runtime. @@ -671,6 +676,16 @@ pub enum CacheStrategy { } impl CacheStrategy { + /// Returns the runtime for CPU-bound SST metadata work for this request. + pub(crate) fn sst_meta_runtime(&self) -> Runtime { + match self { + CacheStrategy::Compaction(_) => common_runtime::compact_runtime(), + CacheStrategy::EnableAll(_) | CacheStrategy::Disabled => { + common_runtime::global_runtime() + } + } + } + /// Returns whether the SST metadata cache is enabled for this strategy. pub(crate) fn sst_meta_cache_enabled(&self) -> bool { match self { @@ -691,7 +706,12 @@ impl CacheStrategy { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { cache_manager - .get_sst_meta_data(file_id, metrics, page_index_policy) + .get_sst_meta_data( + file_id, + metrics, + page_index_policy, + &self.sst_meta_runtime(), + ) .await } CacheStrategy::Disabled => { @@ -878,7 +898,6 @@ impl CacheStrategy { /// Calls [CacheManager::get_range_result()]. /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. - #[allow(dead_code)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, @@ -955,16 +974,6 @@ impl CacheStrategy { } } - /// Calls [CacheManager::vector_index_cache()]. - /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. - #[cfg(feature = "vector_index")] - pub fn vector_index_cache(&self) -> Option<&VectorIndexCacheRef> { - match self { - CacheStrategy::EnableAll(cache_manager) => cache_manager.vector_index_cache(), - CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None, - } - } - /// Calls [CacheManager::puffin_metadata_cache()]. /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. pub fn puffin_metadata_cache(&self) -> Option<&PuffinMetadataCacheRef> { @@ -1023,9 +1032,6 @@ pub struct CacheManager { inverted_index_cache: Option, /// Cache for bloom filter index. bloom_filter_index_cache: Option, - /// Cache for vector index. - #[cfg(feature = "vector_index")] - vector_index_cache: Option, /// Puffin metadata cache. puffin_metadata_cache: Option, /// Cache for time series selectors. @@ -1057,6 +1063,7 @@ impl CacheManager { file_id: RegionFileId, metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Option> { let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id()); let compact = self @@ -1075,7 +1082,7 @@ impl CacheManager { } CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - match compact.decode().await { + match compact.decode(runtime).await { Ok(decoded) => { self.put_sst_meta_data(file_id, decoded.clone()); return Some(decoded); @@ -1101,7 +1108,7 @@ impl CacheManager { let file_cache = write_cache.file_cache(); if self.sst_meta_cache_enabled() { if let Some(metadata) = file_cache - .get_sst_meta_data(key, metrics, page_index_policy) + .get_sst_meta_data(key, metrics, page_index_policy, runtime) .await { metrics.file_cache_hit += 1; @@ -1122,7 +1129,7 @@ impl CacheManager { return Some(decoded); } } else if let Some(decoded) = file_cache - .get_decoded_sst_meta_data(key, metrics, page_index_policy) + .get_decoded_sst_meta_data(key, metrics, page_index_policy, runtime) .await { metrics.file_cache_hit += 1; @@ -1209,7 +1216,7 @@ impl CacheManager { let compact = self .get_compact_sst_meta(&key) .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy))?; - match compact.decode().await { + match compact.decode(&common_runtime::global_runtime()).await { Ok(metadata) => Some(metadata), Err(err) => { warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); @@ -1343,11 +1350,6 @@ impl CacheManager { cache.invalidate_file(file_id.file_id()); } - #[cfg(feature = "vector_index")] - if let Some(cache) = &self.vector_index_cache { - cache.invalidate_file(file_id.file_id()); - } - if let Some(cache) = &self.puffin_metadata_cache { cache.remove(&file_id.to_string()); } @@ -1388,7 +1390,6 @@ impl CacheManager { } /// Gets cached result for range scan. - #[allow(dead_code)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, @@ -1438,11 +1439,6 @@ impl CacheManager { self.bloom_filter_index_cache.as_ref() } - #[cfg(feature = "vector_index")] - pub(crate) fn vector_index_cache(&self) -> Option<&VectorIndexCacheRef> { - self.vector_index_cache.as_ref() - } - pub(crate) fn puffin_metadata_cache(&self) -> Option<&PuffinMetadataCacheRef> { self.puffin_metadata_cache.as_ref() } @@ -1632,9 +1628,6 @@ impl CacheManagerBuilder { self.index_content_size, self.index_content_page_size, ); - #[cfg(feature = "vector_index")] - let vector_index_cache = (self.index_content_size != 0) - .then(|| Arc::new(VectorIndexCache::new(self.index_content_size))); let index_result_cache = (self.index_result_cache_size != 0) .then(|| IndexResultCache::new(self.index_result_cache_size)); let prefilter_result_cache = (self.prefilter_result_cache_size != 0) @@ -1679,8 +1672,6 @@ impl CacheManagerBuilder { write_cache: self.write_cache, inverted_index_cache: Some(Arc::new(inverted_index_cache)), bloom_filter_index_cache: Some(Arc::new(bloom_filter_index_cache)), - #[cfg(feature = "vector_index")] - vector_index_cache, puffin_metadata_cache: Some(Arc::new(puffin_metadata_cache)), selector_result_cache, range_result_cache, @@ -2144,7 +2135,12 @@ mod tests { cache.put_parquet_meta_data(file_id, metadata, None); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime() + ) .await .is_none() ); @@ -2205,14 +2201,24 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime() + ) .await .is_none() ); let (metadata, region_metadata) = sst_parquet_meta(); cache.put_parquet_meta_data(file_id, metadata, None); let cached = cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime(), + ) .await .unwrap(); assert_eq!(region_metadata, cached.region_metadata()); @@ -2230,7 +2236,12 @@ mod tests { cache.remove_parquet_meta_data(file_id); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime() + ) .await .is_none() ); @@ -2247,7 +2258,12 @@ mod tests { cache.put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone())); let cached = cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime(), + ) .await .unwrap(); assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); @@ -2285,9 +2301,15 @@ mod tests { .set_offset_index(Some(offset_indexes)) .build(); let expected_rows = metadata.file_metadata().num_rows(); - let prepared = prepare_sst_meta("test.parquet", metadata, None, PageIndexPolicy::Required) - .await - .unwrap(); + let prepared = prepare_sst_meta( + "test.parquet", + metadata, + None, + PageIndexPolicy::Required, + &common_runtime::global_runtime(), + ) + .await + .unwrap(); let SstMetaPreparation::Prepared(prepared) = prepared else { panic!("valid metadata should produce a compact cache entry"); }; @@ -2296,7 +2318,12 @@ mod tests { cache.put_prepared_sst_meta(file_id, prepared, false); let mut metrics = MetadataCacheMetrics::default(); let cached = cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Required) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Required, + &common_runtime::global_runtime(), + ) .await .unwrap(); @@ -2377,7 +2404,12 @@ mod tests { let mut metrics = MetadataCacheMetrics::default(); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Optional) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Optional, + &common_runtime::global_runtime() + ) .await .is_none() ); @@ -2397,7 +2429,12 @@ mod tests { let mut metrics = MetadataCacheMetrics::default(); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Optional) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Optional, + &common_runtime::global_runtime() + ) .await .is_some() ); @@ -2406,7 +2443,12 @@ mod tests { let mut metrics = MetadataCacheMetrics::default(); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Skip) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Skip, + &common_runtime::global_runtime() + ) .await .is_some() ); diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 0594a1ac10d..4cb60ed8d48 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -21,6 +21,7 @@ use std::time::{Duration, Instant}; use bytes::Bytes; use common_base::readable_size::ReadableSize; +use common_runtime::Runtime; use common_telemetry::{debug, error, info, warn}; use futures::{AsyncWriteExt, FutureExt, TryStreamExt}; use moka::future::Cache; @@ -627,12 +628,13 @@ impl FileCache { key: IndexKey, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Option { let file_path = self.inner.cache_file_path(key); let metadata = self .get_parquet_meta_data(key, cache_metrics, page_index_policy) .await?; - match prepare_sst_meta(&file_path, metadata, None, page_index_policy).await { + match prepare_sst_meta(&file_path, metadata, None, page_index_policy, runtime).await { Ok(metadata) => Some(metadata), Err(err) => { CACHE_MISS @@ -653,12 +655,13 @@ impl FileCache { key: IndexKey, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Option> { let file_path = self.inner.cache_file_path(key); let metadata = self .get_parquet_meta_data(key, cache_metrics, page_index_policy) .await?; - match decode_sst_meta(&file_path, metadata, None, page_index_policy).await { + match decode_sst_meta(&file_path, metadata, None, page_index_policy, runtime).await { Ok(metadata) => Some(metadata), Err(err) => { CACHE_MISS diff --git a/src/mito2/src/cache/index.rs b/src/mito2/src/cache/index.rs index a4049787149..26f9a3d5f96 100644 --- a/src/mito2/src/cache/index.rs +++ b/src/mito2/src/cache/index.rs @@ -15,9 +15,8 @@ pub mod bloom_filter_index; pub mod inverted_index; pub mod result_cache; -#[cfg(feature = "vector_index")] -pub mod vector_index; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::hash::Hash; use std::ops::Range; @@ -46,16 +45,6 @@ pub struct IndexCacheMetrics { pub page_bytes: u64, } -impl IndexCacheMetrics { - /// Merges another set of metrics into this one. - pub fn merge(&mut self, other: &Self) { - self.cache_hit += other.cache_hit; - self.cache_miss += other.cache_miss; - self.num_pages += other.num_pages; - self.page_bytes += other.page_bytes; - } -} - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct PageKey { page_id: u64, @@ -245,6 +234,85 @@ where )) } + /// Like [`IndexCache::get_or_load`] for several ranges, but loads the missing pages of + /// all ranges with a single `load` call so a cold read costs one round trip instead of + /// one per range. + async fn get_or_load_vec( + &self, + key: K, + file_size: u64, + ranges: &[Range], + load: F, + ) -> Result<(Vec, IndexCacheMetrics), E> + where + F: FnOnce(Vec>) -> Fut, + Fut: Future, E>>, + E: std::error::Error, + { + let mut metrics = IndexCacheMetrics::default(); + let mut pages: HashMap = HashMap::new(); + let mut missing: Vec = Vec::new(); + let mut seen = HashSet::new(); + for range in ranges.iter().filter(|r| r.end > r.start) { + let size = (range.end - range.start) as u32; + for page_key in PageKey::generate_page_keys(range.start, size, self.page_size) { + if !seen.insert(page_key.page_id) { + continue; + } + metrics.num_pages += 1; + match self.get_page(key, page_key) { + Some(page) => { + CACHE_HIT.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); + metrics.cache_hit += 1; + metrics.page_bytes += page.len() as u64; + pages.insert(page_key.page_id, page); + } + None => { + CACHE_MISS.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); + metrics.cache_miss += 1; + missing.push(page_key.page_id); + } + } + } + } + + if !missing.is_empty() { + missing.sort_unstable(); + let load_ranges = missing + .iter() + .map(|page_id| { + let start = page_id * self.page_size; + start..(start + self.page_size).min(file_size) + }) + .collect(); + let loaded = load(load_ranges).await?; + debug_assert_eq!(loaded.len(), missing.len()); + for (page_id, page) in missing.into_iter().zip(loaded) { + metrics.page_bytes += page.len() as u64; + self.put_page(key, PageKey { page_id }, page.clone()); + pages.insert(page_id, page); + } + } + + let result = ranges + .iter() + .map(|range| { + let size = (range.end - range.start) as u32; + if size == 0 { + return Bytes::new(); + } + let buffer = Buffer::from_iter( + PageKey::generate_page_keys(range.start, size, self.page_size) + .map(|p| pages[&p.page_id].clone()), + ); + buffer + .slice(PageKey::calculate_range(range.start, size, self.page_size)) + .to_bytes() + }) + .collect(); + Ok((result, metrics)) + } + fn get_page(&self, key: K, page_key: PageKey) -> Option { self.index.get(&(key, page_key)) } diff --git a/src/mito2/src/cache/index/bloom_filter_index.rs b/src/mito2/src/cache/index/bloom_filter_index.rs index 707afb18a65..57f0103c1bb 100644 --- a/src/mito2/src/cache/index/bloom_filter_index.rs +++ b/src/mito2/src/cache/index/bloom_filter_index.rs @@ -154,25 +154,16 @@ impl BloomFilterReader for CachedBloomFilterIndexBl ) -> Result> { let start = metrics.as_ref().map(|_| Instant::now()); - let mut pages = Vec::with_capacity(ranges.len()); - let mut total_cache_metrics = crate::cache::index::IndexCacheMetrics::default(); - for range in ranges { - let inner = &self.inner; - let (page, cache_metrics) = self - .cache - .get_or_load( - (self.file_id, self.index_version, self.column_id, self.tag), - self.blob_size, - range.start, - (range.end - range.start) as u32, - move |ranges| async move { inner.read_vec(&ranges, None).await }, - ) - .await?; - - total_cache_metrics.merge(&cache_metrics); - pages.push(Bytes::from(page)); - } - + let inner = &self.inner; + let (pages, total_cache_metrics) = self + .cache + .get_or_load_vec( + (self.file_id, self.index_version, self.column_id, self.tag), + self.blob_size, + ranges, + move |ranges| async move { inner.read_vec(&ranges, None).await }, + ) + .await?; if let Some(m) = metrics { m.total_ranges += total_cache_metrics.num_pages; m.total_bytes += total_cache_metrics.page_bytes; @@ -190,7 +181,7 @@ impl BloomFilterReader for CachedBloomFilterIndexBl async fn metadata( &self, metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result { + ) -> Result> { if let Some(cached) = self.cache .get_metadata((self.file_id, self.index_version, self.column_id, self.tag)) @@ -199,12 +190,12 @@ impl BloomFilterReader for CachedBloomFilterIndexBl if let Some(m) = metrics { m.cache_hit += 1; } - Ok((*cached).clone()) + Ok(cached) } else { let meta = self.inner.metadata(metrics).await?; self.cache.put_metadata( (self.file_id, self.index_version, self.column_id, self.tag), - Arc::new(meta.clone()), + meta.clone(), ); CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc(); Ok(meta) diff --git a/src/mito2/src/cache/index/inverted_index.rs b/src/mito2/src/cache/index/inverted_index.rs index 4c1b07126c8..675ed736110 100644 --- a/src/mito2/src/cache/index/inverted_index.rs +++ b/src/mito2/src/cache/index/inverted_index.rs @@ -130,25 +130,16 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead ) -> Result> { let start = metrics.as_ref().map(|_| Instant::now()); - let mut pages = Vec::with_capacity(ranges.len()); - let mut total_cache_metrics = crate::cache::index::IndexCacheMetrics::default(); - for range in ranges { - let inner = &self.inner; - let (page, cache_metrics) = self - .cache - .get_or_load( - (self.file_id, self.index_version), - self.blob_size, - range.start, - (range.end - range.start) as u32, - move |ranges| async move { inner.read_vec(&ranges, None).await }, - ) - .await?; - - total_cache_metrics.merge(&cache_metrics); - pages.push(Bytes::from(page)); - } - + let inner = &self.inner; + let (pages, total_cache_metrics) = self + .cache + .get_or_load_vec( + (self.file_id, self.index_version), + self.blob_size, + ranges, + move |ranges| async move { inner.read_vec(&ranges, None).await }, + ) + .await?; if let Some(m) = metrics { m.total_bytes += total_cache_metrics.page_bytes; m.total_ranges += total_cache_metrics.num_pages; @@ -472,4 +463,65 @@ mod test { assert_eq!(read, expected); } } + + #[tokio::test] + async fn test_get_or_load_vec_loads_missing_pages_once() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let mut rng = rand::rng(); + let mut data = vec![0u8; 64 * 1024]; + rng.fill_bytes(&mut data); + let file_size = data.len() as u64; + let cache = InvertedIndexCache::new(1024 * 1024, 1024 * 1024, 1000); + let key = (FileId::random(), 0); + + for _ in 0..FUZZ_REPEAT_TIMES { + cache.invalidate_file(key.0); + let ranges = (0..rng.random_range(1..20)) + .map(|_| { + let start = rng.random_range(0..file_size); + start..rng.random_range(start..=file_size.min(start + 5000)) + }) + .collect::>(); + let expected = ranges + .iter() + .map(|r| bytes::Bytes::copy_from_slice(&data[r.start as usize..r.end as usize])) + .collect::>(); + + let loads = AtomicUsize::new(0); + let load = |pages: Vec>| { + loads.fetch_add(1, Ordering::Relaxed); + // Pages are requested once each, in file order. + assert!(pages.windows(2).all(|w| w[0].end <= w[1].start)); + let data = &data; + async move { + Ok::<_, std::io::Error>( + pages + .iter() + .map(|r| { + bytes::Bytes::copy_from_slice( + &data[r.start as usize..r.end as usize], + ) + }) + .collect(), + ) + } + }; + let (cold, cold_metrics) = cache + .get_or_load_vec(key, file_size, &ranges, load) + .await + .unwrap(); + assert_eq!(cold, expected); + let cold_loads = usize::from(cold_metrics.cache_miss > 0); + assert_eq!(loads.load(Ordering::Relaxed), cold_loads); + + let (warm, metrics) = cache + .get_or_load_vec(key, file_size, &ranges, load) + .await + .unwrap(); + assert_eq!(warm, expected); + assert_eq!(metrics.cache_miss, 0); + assert_eq!(loads.load(Ordering::Relaxed), cold_loads); + } + } } diff --git a/src/mito2/src/cache/index/vector_index.rs b/src/mito2/src/cache/index/vector_index.rs deleted file mode 100644 index 222c8bb0690..00000000000 --- a/src/mito2/src/cache/index/vector_index.rs +++ /dev/null @@ -1,134 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::sync::Arc; - -use index::vector::apply::IndexApplier; -use moka::notification::RemovalCause; -use moka::sync::Cache; -use store_api::storage::{ColumnId, FileId, IndexVersion}; - -use crate::metrics::{CACHE_BYTES, CACHE_EVICTION}; - -const VECTOR_INDEX_CACHE_TYPE: &str = "vector_index"; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct VectorIndexCacheKey { - file_id: FileId, - index_version: IndexVersion, - column_id: ColumnId, -} - -impl VectorIndexCacheKey { - pub fn new(file_id: FileId, index_version: IndexVersion, column_id: ColumnId) -> Self { - Self { - file_id, - index_version, - column_id, - } - } -} - -/// Cached vector index wrapping an `IndexApplier` for efficient searches. -pub struct CachedVectorIndex { - /// The underlying vector index applier. - applier: Box, - /// Size in bytes for cache weighing. - size_bytes: usize, -} - -impl CachedVectorIndex { - /// Creates a new `CachedVectorIndex` from an `IndexApplier`. - pub fn new(applier: Box) -> Self { - let size_bytes = applier.memory_usage(); - Self { - applier, - size_bytes, - } - } - - /// Returns the underlying applier for searching. - pub fn applier(&self) -> &dyn IndexApplier { - self.applier.as_ref() - } - - /// Returns the cache size in bytes. - pub fn size_bytes(&self) -> usize { - self.size_bytes - } -} - -impl std::fmt::Debug for CachedVectorIndex { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let null_count = self.applier.total_rows() - self.applier.indexed_rows(); - f.debug_struct("CachedVectorIndex") - .field("size_bytes", &self.size_bytes) - .field("dimensions", &self.applier.dimensions()) - .field("metric", &self.applier.metric()) - .field("total_rows", &self.applier.total_rows()) - .field("indexed_rows", &self.applier.indexed_rows()) - .field("null_count", &null_count) - .finish() - } -} - -pub struct VectorIndexCache { - inner: Cache>, -} - -pub type VectorIndexCacheRef = Arc; - -impl VectorIndexCache { - pub fn new(capacity: u64) -> Self { - fn to_str(cause: RemovalCause) -> &'static str { - match cause { - RemovalCause::Expired => "expired", - RemovalCause::Explicit => "explicit", - RemovalCause::Replaced => "replaced", - RemovalCause::Size => "size", - } - } - - let inner = Cache::builder() - .max_capacity(capacity) - .weigher(|_k, v: &Arc| v.size_bytes() as u32) - .eviction_listener(|_k, v, cause| { - CACHE_BYTES - .with_label_values(&[VECTOR_INDEX_CACHE_TYPE]) - .sub(v.size_bytes() as i64); - CACHE_EVICTION - .with_label_values(&[VECTOR_INDEX_CACHE_TYPE, to_str(cause)]) - .inc(); - }) - .build(); - Self { inner } - } - - pub fn get(&self, key: &VectorIndexCacheKey) -> Option> { - self.inner.get(key) - } - - pub fn insert(&self, key: VectorIndexCacheKey, value: Arc) { - CACHE_BYTES - .with_label_values(&[VECTOR_INDEX_CACHE_TYPE]) - .add(value.size_bytes() as i64); - self.inner.insert(key, value); - } - - pub fn invalidate_file(&self, file_id: FileId) { - let _ = self - .inner - .invalidate_entries_if(move |k, _| k.file_id == file_id); - } -} diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 0e96c2f16a2..afe943eef50 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -266,8 +266,6 @@ impl WriteCache { inverted_index_config: write_request.inverted_index_config, fulltext_index_config: write_request.fulltext_index_config, bloom_filter_index_config: write_request.bloom_filter_index_config, - #[cfg(feature = "vector_index")] - vector_index_config: write_request.vector_index_config, }; let cleaner = TempFileCleaner::new(region_id, store.clone()); @@ -703,7 +701,6 @@ mod tests { op_type, metadata, source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: Default::default(), @@ -713,8 +710,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let upload_request = SstUploadRequest { @@ -790,7 +785,8 @@ mod tests { async fn test_read_metadata_from_write_cache() { common_telemetry::init_default_ut_logging(); let mut env = TestEnv::new().await; - let data_home = env.data_home().display().to_string(); + // Object keys are relative to the mock store root, not native paths. + let data_home = "write-cache-test".to_string(); let mock_store = env.init_object_store_manager(); let local_dir = create_temp_dir(""); @@ -822,7 +818,6 @@ mod tests { op_type: OperationType::Flush, metadata, source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: cache_manager.clone(), @@ -832,8 +827,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let write_opts = WriteOptions { row_group_size: 512, @@ -917,7 +910,6 @@ mod tests { op_type: OperationType::Flush, metadata, source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: cache_manager.clone(), @@ -927,8 +919,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let write_opts = WriteOptions { row_group_size: 512, diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index c7d55fb8a28..1a025bcaeec 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -461,7 +461,6 @@ impl SstMerger for DefaultSstMerger { write_opts: WriteOptions, ) -> Result<(Vec, Vec)> { let region_id = compaction_region.region_id; - let storage = compaction_region.region_options.storage.clone(); let index_options = compaction_region .current_version .options @@ -479,8 +478,6 @@ impl SstMerger for DefaultSstMerger { let inverted_index_config = compaction_region.engine_config.inverted_index.clone(); let fulltext_index_config = compaction_region.engine_config.fulltext_index.clone(); let bloom_filter_index_config = compaction_region.engine_config.bloom_filter_index.clone(); - #[cfg(feature = "vector_index")] - let vector_index_config = compaction_region.engine_config.vector_index.clone(); let input_file_names = output .inputs @@ -510,7 +507,6 @@ impl SstMerger for DefaultSstMerger { metadata: region_metadata.clone(), source, cache_manager: compaction_region.cache_manager.clone(), - storage, // Readers resolve file overrides before merge/dedup. Replacing // their effective sequences here could promote old rows above // versions in SSTs that were not part of this merge. @@ -526,8 +522,6 @@ impl SstMerger for DefaultSstMerger { inverted_index_config, fulltext_index_config, bloom_filter_index_config, - #[cfg(feature = "vector_index")] - vector_index_config, }, &write_opts, &mut metrics, diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs index f381f80722e..cb2a9d34ec8 100644 --- a/src/mito2/src/compaction/json2.rs +++ b/src/mito2/src/compaction/json2.rs @@ -19,7 +19,9 @@ use arrow_schema::extension::ExtensionType; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::{JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, JsonMetadata}; -use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint}; +use datatypes::json::{ + JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint, TypeHintMismatchPolicy, +}; use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::JsonNativeType; use datatypes::vectors::json::array::JsonArray; @@ -260,7 +262,7 @@ fn select_dynamic_hints( .filter(|(path, stat)| { !stat.is_type_conflicted // TODO(LFC): Instead of "primitive only", consider retaining stable compound types - // that are safe to write to Parquet, as flush does. Or better, unite the two + // that are safe to write to Parquet, as flush does. Or better, unite the two // selection process. && stat.data_type.is_primitive() && !has_ancestor_path(path) @@ -327,7 +329,12 @@ pub(crate) fn rewrite_json2_batch( }; let array = JsonArray::from(array) - .rewrite_to_v2(field, &plan.logical_settings, &plan.target_layout) + .rewrite_to_v2_with_type_hint_mismatch_policy( + field, + &plan.logical_settings, + &plan.target_layout, + TypeHintMismatchPolicy::CoerceOrNull, + ) .context(ConvertValueSnafu)?; debug_assert_eq!( &json2_physical_data_type(&plan.target_layout), diff --git a/src/mito2/src/compaction/run.rs b/src/mito2/src/compaction/run.rs index 2ad9f2f1eba..664005b9af3 100644 --- a/src/mito2/src/compaction/run.rs +++ b/src/mito2/src/compaction/run.rs @@ -475,7 +475,11 @@ where for item in items { let (start, _) = item.range(); - while let Some(run) = runs_sort_by_end.pop_if(|x| x.run.end.unwrap() <= start) { + while runs_sort_by_end + .peek() + .is_some_and(|x| x.run.end.unwrap() <= start) + { + let run = runs_sort_by_end.pop().unwrap(); runs_sort_by_index.push(Wrapper(run)); } diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index 5afcc9c37fd..42c1df23f07 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -546,7 +546,7 @@ async fn test_planning_panic_notifies_and_clears_status() { let finished = recv_compaction_pick_finished(&mut rx).await; let CompactionPlanningResult::Error(err) = &finished.result else { - panic!("expected planning error, got {:?}", &finished.result); + panic!("expected planning error, got {:?}", finished.result); }; assert!(err.to_string().contains("planning boom")); let pending_ddls = scheduler diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index dd407b30382..c170845be3d 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -969,6 +969,7 @@ struct Window { start: Timestamp, end: Timestamp, files: Vec, + /// Start of the half-open window, in seconds. time_window: i64, primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } @@ -1022,12 +1023,7 @@ fn assign_to_windows<'a>( continue; } let (_, end) = f.time_range(); - let time_window = end - .convert_to(TimeUnit::Second) - .unwrap() - .value() - .align_to_ceil_by_bucket(time_window_size) - .unwrap_or(i64::MIN); + let time_window = time_window_start(end, time_window_size); match windows.entry(time_window) { Entry::Occupied(mut e) => { @@ -1043,37 +1039,41 @@ fn assign_to_windows<'a>( windows.into_iter().collect() } +/// Uses the same half-open windows as memtables, including timestamps in the first +/// second after a boundary. An unrepresentable window start is clamped to `i64::MIN`. +fn time_window_start(timestamp: Timestamp, time_window_size: i64) -> i64 { + // Conversion to seconds cannot overflow. + timestamp + .convert_to(TimeUnit::Second) + .and_then(|ts| ts.value().align_by_bucket(time_window_size)) + .unwrap_or(i64::MIN) +} + fn time_window_intersects_range( - window_end: i64, + window_start: i64, time_window_size: i64, time_range: &TimestampRange, ) -> bool { - let first_window = match time_range.start() { - None => i64::MIN, - Some(start) => { - let Some(first_window) = start - .convert_to(TimeUnit::Second) - .and_then(|timestamp| timestamp.value().align_to_ceil_by_bucket(time_window_size)) - else { - return false; - }; - first_window - } - }; + if time_range.is_empty() { + return false; + } + let first_window = time_range + .start() + .map_or(i64::MIN, |start| time_window_start(start, time_window_size)); let last_window = match time_range.end() { None => i64::MAX, Some(end) => { - let Some(last_window) = end + // The range end is exclusive; keep the last second that intersects it. + let Some(last_second) = end .convert_to_ceil(TimeUnit::Second) .and_then(|timestamp| timestamp.value().checked_sub(1)) - .and_then(|timestamp| timestamp.align_to_ceil_by_bucket(time_window_size)) else { return false; }; - last_window + time_window_start(Timestamp::new_second(last_second), time_window_size) } }; - (first_window..=last_window).contains(&window_end) + (first_window..=last_window).contains(&window_start) } fn window_has_overlap(this: &Window, windows: &BTreeMap) -> bool { @@ -1091,7 +1091,7 @@ fn window_has_overlap(this: &Window, windows: &BTreeMap) -> bool { } /// Finds the latest active writing window among all files. -/// Returns `None` when there are no files or all files are corrupted. +/// Returns `None` when there are no files. fn find_latest_window_in_seconds<'a>( files: impl Iterator, time_window_size: i64, @@ -1107,9 +1107,7 @@ fn find_latest_window_in_seconds<'a>( latest_timestamp = Some(end); } } - latest_timestamp - .and_then(|ts| ts.convert_to_ceil(TimeUnit::Second)) - .and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size)) + latest_timestamp.map(|ts| time_window_start(ts, time_window_size)) } /// Finds the active window from the file with the highest sequence number. @@ -1119,9 +1117,7 @@ fn find_latest_window_in_seconds<'a>( /// survives the transient state where level 0 is empty right after its files /// were compacted away. Returns `None` when no file carries a sequence. /// -/// The window key follows the same convention as [`assign_to_windows`] -/// (truncate to seconds, then align up), since it is compared against the -/// window keys produced there. +/// The window key is its inclusive start, as in [`assign_to_windows`]. fn find_active_window_by_sequence<'a>( files: impl Iterator, time_window_size: i64, @@ -1129,12 +1125,7 @@ fn find_active_window_by_sequence<'a>( files .filter(|f| f.meta_ref().sequence.is_some()) .max_by_key(|f| f.meta_ref().sequence) - .and_then(|f| { - f.time_range() - .1 - .convert_to(TimeUnit::Second) - .and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size)) - }) + .map(|f| time_window_start(f.time_range().1, time_window_size)) } #[cfg(test)] @@ -1533,7 +1524,7 @@ mod tests { new_file_handle_with_sequence(FileId::random(), 0, 999, 1, 0), new_file_handle_with_sequence(FileId::random(), 2000, 2999, 0, 0), ]; - assert_eq!(Some(3), active_window_of(&files)); + assert_eq!(Some(2), active_window_of(&files)); // No sequence and no L0: None, exactly as before. let files = [new_file_handle_with_sequence( @@ -1549,7 +1540,7 @@ mod tests { #[test] fn test_get_latest_window_in_seconds() { assert_eq!( - Some(1), + Some(0), find_latest_window_in_seconds([new_file_handle(FileId::random(), 0, 999, 0)].iter(), 1) ); assert_eq!( @@ -1561,7 +1552,7 @@ mod tests { ); assert_eq!( - Some(-9223372036854000), + Some(-9223372036857600), find_latest_window_in_seconds( [new_file_handle(FileId::random(), i64::MIN, i64::MIN + 1, 0)].iter(), 3600, @@ -1569,7 +1560,7 @@ mod tests { ); assert_eq!( - (i64::MAX / 10000000 + 1) * 10000, + (i64::MAX / 10000000) * 10000, find_latest_window_in_seconds( [new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0)].iter(), 10000, @@ -1578,7 +1569,7 @@ mod tests { ); assert_eq!( - Some((i64::MAX / 3600000 + 1) * 3600), + Some((i64::MAX / 3600000) * 3600), find_latest_window_in_seconds( [ new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0), @@ -1590,6 +1581,114 @@ mod tests { ); } + #[test] + fn test_window_assignment_and_activity_at_boundaries() { + // Before, at, and after positive/negative boundaries, including the first + // fractional second: all picker paths must agree with half-open partitions. + for (unit, ticks_per_second) in [ + (TimeUnit::Second, 1), + (TimeUnit::Millisecond, 1_000), + (TimeUnit::Microsecond, 1_000_000), + (TimeUnit::Nanosecond, 1_000_000_000), + ] { + for window_size in [1, 3, 86_400] { + for boundary in [-window_size, 0, window_size] { + for (offset, expected_window) in [ + (-1, boundary - window_size), + (0, boundary), + (ticks_per_second - 1, boundary), + (ticks_per_second, boundary + i64::from(window_size == 1)), + ] { + let timestamp = Timestamp::new(boundary * ticks_per_second + offset, unit); + let file = FileHandle::new( + FileMeta { + time_range: (timestamp, timestamp), + sequence: NonZeroU64::new(1), + ..Default::default() + }, + crate::test_util::new_noop_file_purger(), + ); + let files = [file]; + let windows = assign_to_windows(files.iter(), window_size); + assert_eq!( + vec![expected_window], + windows.keys().copied().collect::>(), + "{timestamp:?}, window size {window_size}" + ); + assert_eq!( + Some(expected_window), + find_active_window_by_sequence(files.iter(), window_size) + ); + assert_eq!( + Some(expected_window), + find_latest_window_in_seconds(files.iter(), window_size) + ); + } + } + } + } + } + + #[test] + fn test_window_assignment_at_timestamp_limits() { + for (seconds, expected_window) in [(i64::MIN, i64::MIN), (i64::MAX, i64::MAX - 1)] { + let timestamp = Timestamp::new_second(seconds); + let files = [FileHandle::new( + FileMeta { + time_range: (timestamp, timestamp), + sequence: NonZeroU64::new(1), + ..Default::default() + }, + crate::test_util::new_noop_file_purger(), + )]; + let windows = assign_to_windows(files.iter(), 3); + assert_eq!( + vec![expected_window], + windows.keys().copied().collect::>() + ); + assert_eq!( + Some(expected_window), + find_active_window_by_sequence(files.iter(), 3) + ); + assert_eq!( + Some(expected_window), + find_latest_window_in_seconds(files.iter(), 3) + ); + assert!(time_window_intersects_range( + expected_window, + 3, + &TimestampRange::single(timestamp) + )); + } + } + + #[tokio::test] + async fn test_midnight_file_is_not_compacted_with_previous_day() { + // Regression for #9441: flush already separates the midnight-only SST. + let midnight = 86_400_000; + CompactionPickerTestCase { + window_size: 86_400, + input_files: [ + (midnight - 240_000, midnight - 180_000), + (midnight - 180_000, midnight - 90_000), + (midnight - 90_000, midnight - 30_000), + (midnight, midnight), + ] + .into_iter() + .enumerate() + .map(|(index, (start, end))| { + new_file_handle_with_sequence(FileId::random(), start, end, 0, index as u64 + 1) + }) + .collect(), + expected_outputs: vec![ExpectedOutput { + input_files: vec![0, 1, 2], + output_level: 1, + }], + } + .check() + .await; + } + #[test] fn test_assign_to_windows() { let windows = assign_to_windows( @@ -1619,7 +1718,7 @@ mod tests { assert_eq!( files[0], windows - .get(&0) + .get(&-3) .unwrap() .files() .next() @@ -1630,7 +1729,7 @@ mod tests { assert_eq!( files[1], windows - .get(&3) + .get(&0) .unwrap() .files() .next() @@ -1641,7 +1740,7 @@ mod tests { assert_eq!( files[2], windows - .get(&12) + .get(&9) .unwrap() .files() .next() @@ -1732,6 +1831,7 @@ mod tests { .collect::>(); let windows = assign_to_windows(file_handles.iter(), time_window); + assert_eq!(expected_files.len(), windows.len()); for (expected_window, overlapping, window_files) in expected_files { let actual_window = windows.get(expected_window).unwrap(); @@ -1756,8 +1856,8 @@ mod tests { &[(0, 999), (1000, 1999), (2000, 2999)], 2, &[ - (0, false, vec![(0, 999)]), - (2, false, vec![(1000, 1999), (2000, 2999)]), + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999)]), ], ); @@ -1774,9 +1874,8 @@ mod tests { &[(0, 999), (1000, 1999), (2000, 2999), (3000, 3999)], 2, &[ - (0, false, vec![(0, 999)]), - (2, false, vec![(1000, 1999), (2000, 2999)]), - (4, false, vec![(3000, 3999)]), + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999), (3000, 3999)]), ], ); @@ -1790,9 +1889,8 @@ mod tests { ], 2, &[ - (0, true, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(0, 3999), (3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(0, 3999), (2000, 2999), (3000, 3999)]), ], ); @@ -1806,43 +1904,39 @@ mod tests { ], 2, &[ - (0, false, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(1999, 3999), (3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(1999, 3999), (2000, 2999), (3000, 3999)]), ], ); check_assign_to_windows_with_overlapping( &[ (0, 999), // window 0 - (1000, 1999), // window 2 + (1000, 1999), // window 0 (2000, 2999), // window 2 - (3000, 3999), // window 4 - (2999, 3999), // window 4 + (3000, 3999), // window 2 + (2999, 3999), // window 2 ], 2, &[ - // window 2 overlaps with window 4 - (0, false, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(2999, 3999), (3000, 3999)]), + // Overlap within a window does not imply cross-window overlap. + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999), (2999, 3999), (3000, 3999)]), ], ); check_assign_to_windows_with_overlapping( &[ (0, 999), // window 0 - (1000, 1999), // window 2 + (1000, 1999), // window 0 (2000, 2999), // window 2 - (3000, 3999), // window 4 - (0, 1000), // // window 2 + (3000, 3999), // window 2 + (0, 2000), // window 2 ], 2, &[ - // only window 0 overlaps with window 2. - (0, true, vec![(0, 999)]), - (2, true, vec![(0, 1000), (1000, 1999), (2000, 2999)]), - (4, false, vec![(3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(0, 2000), (2000, 2999), (3000, 3999)]), ], ); } @@ -1862,7 +1956,7 @@ mod tests { new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 500, - 1999, + 2999, 0, 2, 10, @@ -1872,7 +1966,8 @@ mod tests { let windows = assign_to_windows(files.iter(), 2); - let overlapping = window_has_overlap(windows.get(&2).unwrap(), &windows); + assert_eq!(2, windows.len()); + let overlapping = window_has_overlap(windows.get(&0).unwrap(), &windows); assert!(!overlapping); } @@ -2592,7 +2687,7 @@ mod tests { }; let output = picker - .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), None) + .build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), None) .await .unwrap(); @@ -2609,27 +2704,60 @@ mod tests { #[test] fn test_filter_time_windows_by_time_range() { - let time_range = TimestampRange::new( - Timestamp::new_millisecond(1_200), - Timestamp::new_millisecond(1_800), - ) - .unwrap(); + // Interior, exclusive boundary ends, fractional boundary ends, negative + // timestamps, and empty ranges must all use the same half-open windows. + for (start, end, expected_windows) in [ + (1_200, 1_800, vec![0]), + (0, 3_000, vec![0]), + (3_000, 6_000, vec![3]), + (2_999, 3_000, vec![0]), + (2_999, 3_001, vec![0, 3]), + (3_000, 3_001, vec![3]), + (3_001, 3_002, vec![3]), + (-3_000, 0, vec![-3]), + (-3_001, -2_999, vec![-6, -3]), + (0, 0, vec![]), + (1_500, 1_500, vec![]), + ] { + let time_range = TimestampRange::with_unit(start, end, TimeUnit::Millisecond).unwrap(); + let selected = [-6, -3, 0, 3, 6] + .into_iter() + .filter(|window| time_window_intersects_range(*window, 3, &time_range)) + .collect::>(); + assert_eq!(expected_windows, selected, "{time_range}"); + } + } - assert!(time_window_intersects_range(3, 3, &time_range)); - assert!(!time_window_intersects_range(9, 3, &time_range)); + #[test] + fn test_filter_time_windows_with_unbounded_and_extreme_ranges() { + for (range, expected_windows) in [ + (TimestampRange::min_to_max(), vec![-3, 0, 3]), + ( + TimestampRange::from_start(Timestamp::new_second(0)), + vec![0, 3], + ), + ( + TimestampRange::until_end(Timestamp::new_second(0), false), + vec![-3], + ), + ] { + let selected = [-3, 0, 3] + .into_iter() + .filter(|window| time_window_intersects_range(*window, 3, &range)) + .collect::>(); + assert_eq!(expected_windows, selected, "{range}"); + } - let boundary_range = - TimestampRange::new(Timestamp::new_second(0), Timestamp::new_second(3)).unwrap(); - assert!(time_window_intersects_range(0, 3, &boundary_range)); - assert!(time_window_intersects_range(3, 3, &boundary_range)); - assert!(!time_window_intersects_range(6, 3, &boundary_range)); - - let overflowing_range = TimestampRange::new( + let upper_range = TimestampRange::new( Timestamp::new_second(i64::MAX - 1), Timestamp::new_second(i64::MAX), ) .unwrap(); - assert!(!time_window_intersects_range(0, 4, &overflowing_range)); + assert!(!time_window_intersects_range(0, 4, &upper_range)); + assert!(time_window_intersects_range(i64::MAX - 3, 4, &upper_range)); + + let lower_range = TimestampRange::until_end(Timestamp::new_second(i64::MIN), false); + assert!(!time_window_intersects_range(i64::MIN, 3, &lower_range)); } #[tokio::test] @@ -2659,7 +2787,7 @@ mod tests { }; let output = picker - .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), Some(3)) + .build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), Some(3)) .await .unwrap(); diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 5164e2fd72b..c0fa74281ca 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -92,6 +92,10 @@ pub struct MitoConfig { /// 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, + /// Approximate series and range index size limit in open regions (default: 5 GiB). + /// Workers share a periodically refreshed estimate and skip maintenance when full. + /// In-flight reconciliation can exceed the limit; closed-region files are not counted. + pub experimental_series_index_max_size: ReadableSize, /// Whether to build and query range indexes when series indexes are enabled (default false). /// Obsolete range-index metadata and files are still cleaned up when disabled. pub experimental_enable_range_index: bool, @@ -188,9 +192,6 @@ pub struct MitoConfig { pub fulltext_index: FulltextIndexConfig, /// Bloom filter index configs. pub bloom_filter_index: BloomFilterConfig, - /// Vector index configs (HNSW). - #[cfg(feature = "vector_index")] - pub vector_index: VectorIndexConfig, /// Minimum time interval between two compactions. /// To align with the old behavior, the default value is 0 (no restrictions). @@ -221,6 +222,7 @@ impl Default for MitoConfig { compress_manifest: false, max_background_index_builds: divide_num_cpus(8), experimental_enable_series_index: false, + experimental_series_index_max_size: ReadableSize::gb(5), experimental_enable_range_index: false, experimental_series_index_maintenance_interval: DEFAULT_SERIES_INDEX_MAINTENANCE_INTERVAL, @@ -257,12 +259,11 @@ impl Default for MitoConfig { inverted_index: InvertedIndexConfig::default(), fulltext_index: FulltextIndexConfig::default(), bloom_filter_index: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index: VectorIndexConfig::default(), min_compaction_interval: Duration::from_secs(0), schedule_compaction_after_edit: true, default_flat_format: true, - experimental_series_scan_v2: true, + // FIXME(#9435): Keep v2 opt-in while long-range scan memory usage is investigated. + experimental_series_scan_v2: false, gc: GcConfig::default(), }; @@ -280,6 +281,14 @@ impl MitoConfig { /// /// Returns an error if there is a configuration that unable to sanitize. pub fn sanitize(&mut self, data_home: &str) -> Result<()> { + if self.experimental_enable_series_index { + snafu::ensure!( + self.experimental_series_index_max_size.as_bytes() >= 1024, + crate::error::InvalidConfigSnafu { + reason: "experimental_series_index_max_size must be at least 1KiB" + } + ); + } // Use default value if `num_workers` is 0. if self.num_workers == 0 { self.num_workers = divide_num_cpus(2); @@ -442,12 +451,17 @@ mod tests { let mut config: MitoConfig = toml::from_str( "experimental_enable_series_index = true experimental_enable_range_index = false + experimental_series_index_max_size = '64MiB' experimental_series_index_maintenance_interval = '30s' experimental_series_index_bucket_width = '2days'", ) .unwrap(); config.sanitize("/data").unwrap(); assert!(config.experimental_enable_series_index); + assert_eq!( + ReadableSize::mb(64), + config.experimental_series_index_max_size + ); assert!(!config.experimental_enable_range_index); assert_eq!( config.experimental_series_index_maintenance_interval, @@ -459,6 +473,10 @@ mod tests { ); let restored: MitoConfig = toml::from_str(&toml::to_string(&config).unwrap()).unwrap(); assert_eq!(config, restored); + config.experimental_series_index_max_size = ReadableSize(1023); + assert!(config.sanitize("/data").is_err()); + config.experimental_series_index_max_size = ReadableSize(1024); + config.sanitize("/data").unwrap(); let mut config: MitoConfig = toml::from_str("experimental_series_index_maintenance_interval = '0s'").unwrap(); @@ -481,14 +499,6 @@ mod tests { assert_eq!(config, restored); } } - - #[test] - fn test_experimental_series_scan_v2_config() { - assert!(MitoConfig::default().experimental_series_scan_v2); - - let config: MitoConfig = toml::from_str("experimental_series_scan_v2 = false").unwrap(); - assert!(!config.experimental_series_scan_v2); - } } /// Index build mode. @@ -783,51 +793,6 @@ impl BloomFilterConfig { } } -/// Configuration options for the vector index (HNSW). -#[cfg(feature = "vector_index")] -#[serde_as] -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] -#[serde(default)] -pub struct VectorIndexConfig { - /// Whether to create the index on flush: automatically or never. - pub create_on_flush: Mode, - /// Whether to create the index on compaction: automatically or never. - pub create_on_compaction: Mode, - /// Whether to apply the index on query: automatically or never. - pub apply_on_query: Mode, - /// Memory threshold for creating the index. - pub mem_threshold_on_create: MemoryThreshold, -} - -#[cfg(feature = "vector_index")] -impl Default for VectorIndexConfig { - fn default() -> Self { - Self { - create_on_flush: Mode::Auto, - create_on_compaction: Mode::Auto, - apply_on_query: Mode::Auto, - mem_threshold_on_create: MemoryThreshold::Auto, - } - } -} - -#[cfg(feature = "vector_index")] -impl VectorIndexConfig { - pub fn mem_threshold_on_create(&self) -> Option { - match self.mem_threshold_on_create { - MemoryThreshold::Auto => { - if let Some(sys_memory) = get_total_memory_readable() { - Some((sys_memory / INDEX_CREATE_MEM_THRESHOLD_FACTOR).as_bytes() as usize) - } else { - Some(ReadableSize::mb(64).as_bytes() as usize) - } - } - MemoryThreshold::Unlimited => None, - MemoryThreshold::Size(size) => Some(size.as_bytes() as usize), - } - } -} - /// Divide cpu num by a non-zero `divisor` and returns at least 1. fn divide_num_cpus(divisor: usize) -> usize { debug_assert!(divisor > 0); diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index fbfd339cff7..47752adbe75 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -51,6 +51,10 @@ pub mod listener; #[cfg(test)] mod merge_mode_test; #[cfg(test)] +mod object_store_wal_recovery_test; +#[cfg(test)] +mod object_store_wal_test; +#[cfg(test)] mod open_test; #[cfg(test)] mod parallel_test; diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index fd74c4b498a..cf039b52278 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -37,7 +37,7 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions}; use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::provider::Provider; -use store_api::metadata::ColumnMetadata; +use store_api::metadata::{ColumnMetadata, RegionMetadata}; use store_api::metric_engine_consts::{PRIMARY_KEY_ENCODING, TABLE_COLUMN_METADATA_EXTENSION_KEY}; use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ @@ -193,6 +193,133 @@ fn assert_column_metadatas(column_name: &[(&str, ColumnId)], column_metadatas: & } } +#[tokio::test] +async fn test_sync_columns_preserves_primary_key_order() { + test_sync_columns_preserves_primary_key_order_with_format(false).await; + test_sync_columns_preserves_primary_key_order_with_format(true).await; +} + +async fn assert_sync_columns_readable( + engine: &MitoEngine, + region_id: RegionId, + metadata: &RegionMetadata, + expected: &str, +) { + let current = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(current.primary_key, metadata.primary_key); + let projection = metadata + .column_metadatas + .iter() + .map(|column| current.column_index_by_id(column.column_id).unwrap()) + .collect(); + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + projection: Some(projection), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap(), + expected + ); +} + +async fn test_sync_columns_preserves_primary_key_order_with_format(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 table_dir = request.table_dir.clone(); + let mut row_schema = rows_schema(&request); + row_schema.push(tag_column_schema("tag_1", ColumnDataType::String)); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + engine + .handle_request(region_id, RegionRequest::Alter(add_tag1())) + .await + .unwrap(); + let metadata = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(metadata.primary_key, vec![0, 3]); + + put_rows( + &engine, + region_id, + Rows { + schema: row_schema, + rows: build_rows_for_tags("a", "b", 0, 2, 10), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let before = RecordBatches::try_collect( + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .scan() + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + + let mut columns = metadata.column_metadatas.clone(); + columns.push(ColumnMetadata { + column_schema: ColumnSchema::new("field_1", ConcreteDataType::float64_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 4, + }); + let tags = metadata.primary_key_columns().cloned().collect::>(); + let mut tags = tags.into_iter(); + for column in &mut columns { + if column.semantic_type == SemanticType::Tag { + *column = tags.next().unwrap(); + } + } + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SyncColumns { + column_metadatas: columns, + }, + }), + ) + .await + .unwrap(); + + assert_sync_columns_readable(&engine, region_id, &metadata, &before).await; + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert_sync_columns_readable(&engine, region_id, &metadata, &before).await; +} + #[tokio::test] async fn test_alter_region() { test_alter_region_with_format(false).await; diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 776e1f78abe..35f15f2393c 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -1500,7 +1500,10 @@ async fn check_request_skip_wal_recovery(flat_format: bool, skip_wal: bool) { 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; + // Stop the engine directly, without region close, to exercise WAL-only recovery. + let engine = + restart_without_region_close(&mut env, engine, flat_format, region_id, table_dir).await; + let stream = engine .scan_to_stream(region_id, ScanRequest::default()) .await @@ -1625,7 +1628,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( skip_wal: bool, flush_before_reopen: bool, ) { - use crate::region_write_ctx::RegionWriteCtx; + use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::OptionOutputTx; use crate::test_util::LogStoreImpl; use crate::wal::Wal; @@ -1657,6 +1660,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( ®ion.version_control, region.provider.clone(), None, + WriteSource::Request, ); let mut receivers = Vec::with_capacity(4); for (index, (skip, timestamp)) in [(false, 0), (skip_wal, 0), (false, 0), (skip_wal, 1)] @@ -1692,8 +1696,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( 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(); + ctx.write_memtables().await; drop(ctx); for rx in receivers { assert_eq!(rx.await.unwrap().unwrap(), 2); @@ -1751,9 +1754,13 @@ async fn check_request_skip_wal_mixed_batch_recovery( ); } - // 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; + // Bypass region close so only WAL-backed rows recover unless the explicit + // flush above persisted all requests. + let engine = + restart_without_region_close(&mut env, engine, flat_format, region_id, table_dir).await; + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); 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)] @@ -1830,3 +1837,38 @@ async fn request_skip_wal_values(engine: &MitoEngine, region_id: RegionId) -> Ve values.sort_unstable_by_key(|(timestamp, _)| *timestamp); values } + +/// Restarts workers without closing regions, preserving WAL-only recovery coverage. +async fn restart_without_region_close( + env: &mut TestEnv, + engine: MitoEngine, + flat_format: bool, + region_id: RegionId, + table_dir: String, +) -> MitoEngine { + let engine = env + .reopen_engine( + engine, + MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + options: HashMap::new(), + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine +} diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index ce950f8a246..6d2241846b4 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -563,7 +563,7 @@ async fn run_index_publication_compaction_race(phase: IndexPublicationPhase, gc_ engine_for_build .handle_request( region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await }); @@ -705,7 +705,7 @@ async fn test_index_build_uses_physical_file_region_and_logical_manifest_region( engine .handle_request( target_region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await .unwrap(); @@ -1391,7 +1391,7 @@ async fn test_index_build_type_manual_basic() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 0); // Trigger manual index build task and make sure index file is built without flush or compaction. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); listener.wait_finish(1).await; let scanner = engine @@ -1403,7 +1403,7 @@ async fn test_index_build_type_manual_basic() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); // Test idempotency: Second manual index build request on the same file. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; let scanner = engine @@ -1416,7 +1416,7 @@ async fn test_index_build_type_manual_basic() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); // Test idempotency again: Third manual index build request to further verify. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; let scanner = engine @@ -1475,7 +1475,7 @@ async fn test_index_build_type_manual_consistency() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); // Check index build task for consistent file will be skipped. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); // Reopen the region to ensure the task wasn't skipped due to insufficient time. reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; @@ -1505,53 +1505,13 @@ async fn test_index_build_type_manual_consistency() { // SyncColumns won't trigger index build. assert_listener_counts(&listener, 1, 1); - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); listener.wait_finish(2).await; // previous 1 + new 1 // Because the file is inconsistent, new index build task is triggered. assert_listener_counts(&listener, 2, 2); } -#[tokio::test] -async fn test_gate_index_build_listener_smoke() { - use store_api::storage::{FileId, RegionId}; - - use crate::engine::listener::{EventListener, GateIndexBuildListener}; - use crate::sst::file::RegionFileId; - - let gate = Arc::new(GateIndexBuildListener::default()); - - // Initial counts are zero. - assert_eq!(gate.begin_count(), 0); - assert_eq!(gate.finish_count(), 0); - assert_eq!(gate.abort_count(), 0); - - // Spawn a task that will block in on_index_build_begin. - let gate_clone = gate.clone(); - let handle = tokio::spawn(async move { - gate_clone - .on_index_build_begin(RegionFileId::new(RegionId::new(1, 1), FileId::random())) - .await; - }); - - // Wait for begin to arrive. - tokio::time::timeout(std::time::Duration::from_secs(5), gate.wait_begin(1)) - .await - .unwrap(); - assert_eq!(gate.begin_count(), 1); - assert_eq!(gate.finish_count(), 0); - assert_eq!(gate.abort_count(), 0); - - // Release the blocked begin. - gate.release_begin(); - - // The spawned task should now complete. - tokio::time::timeout(std::time::Duration::from_secs(5), handle) - .await - .unwrap() - .unwrap(); -} - #[tokio::test] async fn test_index_build_type_manual_duplicate_in_flight() { let mut env = TestEnv::with_prefix("test_index_build_type_manual_duplicate_in_flight_").await; @@ -1621,7 +1581,7 @@ async fn test_index_build_type_manual_duplicate_in_flight() { // so we must spawn it in a separate task. let engine_clone = engine.clone(); let first_handle = tokio::spawn(async move { - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine_clone.handle_request(region_id, request).await }); @@ -1636,7 +1596,7 @@ async fn test_index_build_type_manual_duplicate_in_flight() { // Issue the second manual BuildIndex for the same region/file. // Since the file is already in building_files (from the first manual build), // schedule_build detects the duplicate and calls on_index_build_abort. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); // The second request should have been aborted as duplicate. @@ -1719,7 +1679,7 @@ async fn test_reopen_waits_for_active_index_build_of_previous_incarnation() { engine_for_old_build .handle_request( region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await }); @@ -1747,7 +1707,7 @@ async fn test_reopen_waits_for_active_index_build_of_previous_incarnation() { engine_for_new_build .handle_request( region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await }); @@ -1919,3 +1879,38 @@ async fn test_index_build_type_compact_abort_race() { assert_eq!(scanner.num_files(), 1); assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); } + +#[tokio::test] +async fn test_build_series_index_rejects_non_metric_region() { + let mut env = TestEnv::with_prefix("series-unsupported-schema").await; + let engine = env + .create_engine(MitoConfig { + experimental_enable_series_index: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let error = engine + .handle_request( + region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest { + options: Some(api::v1::region::build_index_request::Options::SeriesIndex( + Default::default(), + )), + }), + ) + .await + .unwrap_err(); + assert!( + error.to_string().contains("sparse metric metadata"), + "{error}" + ); + engine.stop().await.unwrap(); +} diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index 6ef4a293a2a..432e3367603 100644 --- a/src/mito2/src/engine/listener.rs +++ b/src/mito2/src/engine/listener.rs @@ -583,12 +583,12 @@ impl EventListener for GateIndexBuildListener { info!("Region {} index build begin (gated)", region_file_id); self.begin_count.fetch_add(1, Ordering::Relaxed); self.begin_notify.notify_one(); - // Block until the test releases the gate. - let _permit = self - .begin_blocker + // Consume the release so subsequent builds remain blocked. + self.begin_blocker .acquire() .await - .expect("gate semaphore should not be closed"); + .expect("gate semaphore should not be closed") + .forget(); } async fn on_index_build_finish(&self, region_file_id: RegionFileId) { diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs new file mode 100644 index 00000000000..126e38e4b26 --- /dev/null +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -0,0 +1,1357 @@ +// 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. + +//! Recovery tests for regions on a real [`ObjectStoreLogStore`] over an +//! in-memory object store. The store seals objects only through its testing +//! hooks, so every test controls which entries are durable. +//! +//! Every open of the store writes a start object above every present object, +//! so the first data object of a store opened on an empty prefix has +//! sequence 1. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use api::v1::Rows; +use common_base::readable_size::ReadableSize; +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; +use common_recordbatch::RecordBatches; +use common_wal::config::object_store::{ + AckMode, CorruptedSegmentAction, ObjectStoreWalConfig, STANDALONE_GENERATION, +}; +use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use log_store::ObjectStoreLogStore; +use log_store::object_store_wal::{WalHole, entry_id}; +use object_store::ObjectStore; +use object_store::services::Memory; +use rstest::rstest; +use store_api::logstore::LogStore; +use store_api::logstore::provider::Provider; +use store_api::mito_engine_options::SKIP_WAL_KEY; +use store_api::region_engine::{RegionEngine, RegionRole}; +use store_api::region_request::{ + PathType, RegionDropRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest, + RegionRequest, RegionTruncateRequest, +}; +use store_api::storage::{RegionId, ScanRequest}; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::region::MitoRegionRef; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, flush_region, put_rows, + rows_schema, +}; + +/// The root the store derives its node prefix from. +const ROOT: &str = "cluster-a/wal"; +/// The node prefix of datanode 0 in the standalone generation under [`ROOT`]. +const PREFIX: &str = "cluster-a/wal/datanodes/0/epochs/0"; +/// The engine runs two workers and these regions map to different ones, so +/// their writes can be admitted into the same open batch. +const REGION_A: RegionId = RegionId::new(1, 1); +const REGION_B: RegionId = RegionId::new(1, 2); +const WAIT: Duration = Duration::from_secs(30); + +fn memory_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap() +} + +/// The configuration of a store with `ack_mode` that never seals a batch on +/// its own. +fn store_config(ack_mode: AckMode) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + storage_provider: String::new(), + prefix: ROOT.to_string(), + flush_interval: Duration::from_secs(3600), + max_batch_bytes: ReadableSize(u64::MAX), + ack_mode, + ..Default::default() + } +} + +/// Opens a store under [`PREFIX`] with `ack_mode` that never seals a batch +/// on its own. +async fn open_store(object_store: &ObjectStore, ack_mode: AckMode) -> Arc { + ObjectStoreLogStore::try_new( + object_store.clone(), + &store_config(ack_mode), + 0, + STANDALONE_GENERATION, + ) + .await + .unwrap() +} + +async fn new_engine(env: &mut TestEnv, store: Arc) -> MitoEngine { + let config = MitoConfig { + num_workers: 2, + ..Default::default() + }; + env.create_engine_with_log_store(config, store).await +} + +fn wal_options() -> HashMap { + let options = WalOptions::ObjectStore(ObjectStoreWalOptions::new(PREFIX.to_string())); + HashMap::from([( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&options).unwrap(), + )]) +} + +fn provider(region_id: RegionId) -> Provider { + Provider::object_store_provider(region_id, PREFIX.to_string()) +} + +/// Creates a region on the object store WAL with `extra_options` and returns +/// its table dir and row schema. +async fn create_region( + engine: &MitoEngine, + region_id: RegionId, + extra_options: &[(&str, &str)], +) -> (String, Vec) { + let mut builder = + CreateRequestBuilder::new().insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]); + for (key, value) in extra_options { + builder = builder.insert_option(key, value); + } + let request = builder.build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + (table_dir, schema) +} + +/// Opens a region whose options select the object store WAL under +/// [`PREFIX`] plus `extra_options`, and makes it writable. +async fn open_region( + engine: &MitoEngine, + region_id: RegionId, + table_dir: &str, + extra_options: &[(&str, &str)], +) -> std::result::Result<(), BoxedError> { + let mut options = wal_options(); + options.extend( + extra_options + .iter() + .map(|(key, value)| (key.to_string(), value.to_string())), + ); + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: table_dir.to_string(), + options, + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await?; + engine.set_region_role(region_id, RegionRole::Leader) +} + +fn rows(schema: &[api::v1::ColumnSchema], start: usize, end: usize) -> Rows { + Rows { + schema: schema.to_vec(), + rows: build_rows(start, end), + } +} + +/// Returns every row of the region as a table, so two scans can be compared +/// row by row. +async fn scan_rows(engine: &MitoEngine, region_id: RegionId) -> String { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap() +} + +fn region(engine: &MitoEngine, region_id: RegionId) -> MitoRegionRef { + engine.get_region(region_id).unwrap() +} + +/// The entry ids a region tracks in memory and in its manifest, and the +/// rows its memtables hold. +#[derive(Debug, PartialEq, Eq)] +struct RegionRecoveryState { + flushed_entry_id: u64, + last_entry_id: u64, + topic_latest_entry_id: u64, + manifest_flushed_entry_id: u64, + memtable_rows: u64, +} + +async fn region_recovery_state(engine: &MitoEngine, region_id: RegionId) -> RegionRecoveryState { + let region = region(engine, region_id); + let current = region.version_control.current(); + let manifest = region.manifest_ctx.manifest().await; + RegionRecoveryState { + flushed_entry_id: current.version.flushed_entry_id, + last_entry_id: current.last_entry_id, + topic_latest_entry_id: region.topic_latest_entry_id.load(Ordering::Relaxed), + manifest_flushed_entry_id: manifest.flushed_entry_id, + memtable_rows: current.version.memtables.num_rows(), + } +} + +const EMPTY_RECOVERY_STATE: RegionRecoveryState = RegionRecoveryState { + flushed_entry_id: 0, + last_entry_id: 0, + topic_latest_entry_id: 0, + manifest_flushed_entry_id: 0, + memtable_rows: 0, +}; + +/// Writes rows through the engine and seals them into one WAL object. +/// +/// A put in the `durable` mode blocks until its entries are durable, so the +/// writes run in the background until the store has admitted all of them and +/// the batch is sealed by hand. +struct SealedWriter { + store: Arc, + admitted: usize, +} + +impl SealedWriter { + fn new(store: &Arc) -> Self { + Self { + store: store.clone(), + admitted: 0, + } + } + + async fn put_and_seal(&mut self, engine: &MitoEngine, writes: Vec<(RegionId, Rows)>) { + let handles = writes + .into_iter() + .map(|(region_id, rows)| { + let engine = engine.clone(); + tokio::spawn(async move { put_rows(&engine, region_id, rows).await }) + }) + .collect::>(); + self.admitted += handles.len(); + tokio::time::timeout(WAIT, self.store.wait_for_admitted_appends(self.admitted)) + .await + .expect("writes must reach the store") + .unwrap(); + self.store.seal_open_batch().await.unwrap(); + for handle in handles { + handle.await.unwrap(); + } + } +} + +fn latest(store: &ObjectStoreLogStore, region_id: RegionId) -> u64 { + store.latest_entry_id(&provider(region_id)).unwrap() +} + +/// Returns the sequences of the WAL objects under the prefix, in order. +async fn wal_object_seqs(object_store: &ObjectStore) -> Vec { + let mut seqs = object_store + .list(&format!("{PREFIX}/objects/")) + .await + .unwrap() + .into_iter() + .filter(|entry| entry.metadata().is_file()) + .map(|entry| { + let path = entry.path(); + path.rsplit('/') + .next() + .and_then(|name| name.strip_suffix(".wal")) + .and_then(|seq| seq.parse().ok()) + .unwrap_or_else(|| panic!("unexpected WAL object key {path}")) + }) + .collect::>(); + seqs.sort_unstable(); + seqs +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_reopen_after_partial_flush_replays_only_unflushed_regions(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-partial-flush").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + + // Objects 1 and 2 each hold one entry of both regions: the ids of a + // region are its position under the sequence of its object. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 0, 2)), + (REGION_B, rows(&schema_b, 0, 3)), + ], + ) + .await; + writer + .put_and_seal( + &engine, + vec![ + (REGION_B, rows(&schema_b, 3, 5)), + (REGION_A, rows(&schema_a, 2, 4)), + ], + ) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!(entry_id(2, 1), latest(&store, REGION_B)); + + // Flushing empties the memtables, so the topic latest entry id follows + // the store. + flush_region(&engine, REGION_A, None).await; + let flushed_a = RegionRecoveryState { + flushed_entry_id: entry_id(2, 1), + last_entry_id: entry_id(2, 1), + topic_latest_entry_id: entry_id(2, 1), + manifest_flushed_entry_id: entry_id(2, 1), + memtable_rows: 0, + }; + assert_eq!(flushed_a, region_recovery_state(&engine, REGION_A).await); + let rows_a = scan_rows(&engine, REGION_A).await; + let rows_b = scan_rows(&engine, REGION_B).await; + + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap(); + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + + // Region A replays nothing: it flushed both entries, so the topic latest + // entry id comes from the store. Region B replays both of its entries. + assert_eq!(flushed_a, region_recovery_state(&engine, REGION_A).await); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 5, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_B).await + ); + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!(entry_id(2, 1), latest(&store, REGION_B)); + assert_eq!(rows_a, scan_rows(&engine, REGION_A).await); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + assert_eq!(4, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!(5, engine.get_region_statistic(REGION_B).unwrap().num_rows); +} + +/// Waits until no WAL object delete is in flight. +async fn wait_for_collection(store: &ObjectStoreLogStore) { + tokio::time::timeout(WAIT, store.wait_for_garbage_collection()) + .await + .expect("collection must complete"); +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_flush_collects_objects_below_the_watermark(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-collection").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + + // Object 1 holds entries of both regions, object 2 of region A and + // object 3 of region B. Object 1 carries the boundary past the start + // object, which is collected. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 0, 2)), + (REGION_B, rows(&schema_b, 0, 3)), + ], + ) + .await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 2, 4))]) + .await; + writer + .put_and_seal(&engine, vec![(REGION_B, rows(&schema_b, 3, 5))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![1, 2, 3], wal_object_seqs(&object_store).await); + + // Region B has no watermark and still needs object 1, so flushing + // region A collects nothing. + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 4, 6))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![1, 2, 3, 4], wal_object_seqs(&object_store).await); + + // Once region B is flushed, the first object a region needs is object 4, + // the unflushed entry of region A; object 5 carries that boundary. + flush_region(&engine, REGION_B, None).await; + writer + .put_and_seal(&engine, vec![(REGION_B, rows(&schema_b, 5, 7))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![4, 5], wal_object_seqs(&object_store).await); + let rows_a = scan_rows(&engine, REGION_A).await; + let rows_b = scan_rows(&engine, REGION_B).await; + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // Recovery stops at the boundary and each region replays its entry above + // its flushed id. + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap(); + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + for (region_id, flushed_entry_id, last_entry_id) in [ + (REGION_A, entry_id(2, 1), entry_id(4, 1)), + (REGION_B, entry_id(3, 1), entry_id(5, 1)), + ] { + let state = region_recovery_state(&engine, region_id).await; + assert_eq!(flushed_entry_id, state.flushed_entry_id); + assert_eq!(last_entry_id, state.last_entry_id); + assert_eq!(2, state.memtable_rows); + } + assert_eq!(rows_a, scan_rows(&engine, REGION_A).await); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + assert_eq!(6, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!(7, engine.get_region_statistic(REGION_B).unwrap().num_rows); + wait_for_collection(&store).await; + assert_eq!(vec![4, 5, 6], wal_object_seqs(&object_store).await); +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_reopen_after_crash_replays_durable_entries_once(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-crash").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is durable as an object but the manifest never learns of it. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 0, 3))]) + .await; + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + let rows_before = scan_rows(&engine, REGION_A).await; + + drop(writer); + crash(engine, store).await; + + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(1, 1), + memtable_rows: 3, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(3, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[rstest] +#[case(CorruptedSegmentAction::Skip)] +#[case(CorruptedSegmentAction::Fail)] +#[tokio::test] +async fn test_reopen_with_a_corrupted_segment( + #[case] on_corrupted_segment: CorruptedSegmentAction, +) { + let mut env = TestEnv::with_prefix("object-store-wal-corrupted-segment").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + + // Objects 1 and 2 each hold one entry of both regions; nothing is flushed. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 0, 2)), + (REGION_B, rows(&schema_b, 0, 3)), + ], + ) + .await; + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 2, 4)), + (REGION_B, rows(&schema_b, 3, 5)), + ], + ) + .await; + let rows_b = scan_rows(&engine, REGION_B).await; + engine.stop().await.unwrap(); + + // The segment of region A in object 1 is damaged. + let (path, segment) = store + .segment_location(&provider(REGION_A), 1) + .unwrap() + .unwrap(); + let mut bytes = object_store.read(&path).await.unwrap().to_vec(); + bytes[segment.end as usize - 1] ^= 1; + object_store.write(&path, bytes).await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + let config = ObjectStoreWalConfig { + on_corrupted_segment, + ..store_config(AckMode::Durable) + }; + let store = + ObjectStoreLogStore::try_new(object_store.clone(), &config, 0, STANDALONE_GENERATION) + .await + .unwrap(); + let engine = new_engine(&mut env, store.clone()).await; + + // Region B replays both of its entries whatever the damage does to A. + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + assert_eq!(5, engine.get_region_statistic(REGION_B).unwrap().num_rows); + assert!(store.wal_holes(&provider(REGION_B)).unwrap().is_empty()); + + let opened = open_region(&engine, REGION_A, &table_dir_a, &[]).await; + match on_corrupted_segment { + // Region A replays the entry of object 2 alone and records the + // segment of object 1 as a hole. + CorruptedSegmentAction::Skip => { + opened.unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 2, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!( + vec![WalHole { + path, + object_seq: 1, + min_entry_id: entry_id(1, 1), + max_entry_id: entry_id(1, 1), + }], + store.wal_holes(&provider(REGION_A)).unwrap() + ); + } + // The read fails, so the region does not open. + CorruptedSegmentAction::Fail => { + let message = opened.unwrap_err().output_msg(); + assert!( + message.contains("checksum mismatch"), + "unexpected error: {message}" + ); + assert!(!engine.is_region_exists(REGION_A)); + assert!(store.wal_holes(&provider(REGION_A)).unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn test_open_region_rejects_mismatched_wal_prefix() { + let mut env = TestEnv::with_prefix("object-store-wal-prefix-mismatch").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 0, 3))]) + .await; + + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The process now runs its store under the next generation while the + // region still persists the prefix it was created with. + let config = store_config(AckMode::Durable); + let other_prefix = config.node_prefix(0, STANDALONE_GENERATION + 1); + let store = + ObjectStoreLogStore::try_new(object_store.clone(), &config, 0, STANDALONE_GENERATION + 1) + .await + .unwrap(); + let engine = new_engine(&mut env, store).await; + let err = open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + let message = err.output_msg(); + assert!( + message.contains(PREFIX) && message.contains(&other_prefix), + "unexpected error: {message}" + ); + assert!(!engine.is_region_exists(REGION_A)); +} + +#[tokio::test] +async fn test_reopen_on_empty_prefix_without_durable_entries() { + let mut env = TestEnv::with_prefix("object-store-wal-empty-prefix").await; + let store = open_store(&memory_store(), AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + // Region A never writes; region B skips the WAL and flushes its rows. + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[(SKIP_WAL_KEY, "true")]).await; + put_rows(&engine, REGION_B, rows(&schema_b, 0, 3)).await; + flush_region(&engine, REGION_B, None).await; + let rows_b = scan_rows(&engine, REGION_B).await; + + engine.stop().await.unwrap(); + drop(engine); + drop(store); + + // A fresh object store holds no object under the prefix. + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap(); + open_region(&engine, REGION_B, &table_dir_b, &[(SKIP_WAL_KEY, "true")]) + .await + .unwrap(); + + assert_eq!(0, latest(&store, REGION_A)); + assert_eq!(0, latest(&store, REGION_B)); + assert_eq!( + EMPTY_RECOVERY_STATE, + region_recovery_state(&engine, REGION_A).await + ); + // Region B keeps the object store provider but wrote nothing to it; its + // rows come back from the SST alone. + assert_eq!(provider(REGION_B), region(&engine, REGION_B).provider); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_B).await.memtable_rows + ); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + + // The region is usable and its first entry lands in the first object + // after the start object, which carries the boundary past the start + // object. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 0, 2))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![1], wal_object_seqs(&object_store).await); + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_entry_ids_continue_across_two_restarts(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-two-restarts").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry of object 1 is flushed, the entry of object 2 is not. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 0, 2))]) + .await; + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 2, 4))]) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The first restart writes start object 3 and replays the entry of + // object 2 only. + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(1, 1), + last_entry_id: entry_id(2, 1), + topic_latest_entry_id: entry_id(1, 1), + manifest_flushed_entry_id: entry_id(1, 1), + memtable_rows: 2, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(4, engine.get_region_statistic(REGION_A).unwrap().num_rows); + + // The sequence continues at object 4, whose entry is flushed; the entry + // of object 5 is not. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 4, 6))]) + .await; + assert_eq!(entry_id(4, 1), latest(&store, REGION_A)); + assert_eq!( + entry_id(4, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 6, 8))]) + .await; + assert_eq!(entry_id(5, 1), latest(&store, REGION_A)); + let rows_before = scan_rows(&engine, REGION_A).await; + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The second restart replays the entry of object 5 only. + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(4, 1), + last_entry_id: entry_id(5, 1), + topic_latest_entry_id: entry_id(4, 1), + manifest_flushed_entry_id: entry_id(4, 1), + memtable_rows: 2, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(entry_id(5, 1), latest(&store, REGION_A)); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(8, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!( + 8, + region(&engine, REGION_A) + .version_control + .committed_sequence() + ); +} + +/// Sends a put without asserting its result. +async fn try_put( + engine: &MitoEngine, + region_id: RegionId, + rows: Rows, +) -> std::result::Result<(), BoxedError> { + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows, + hint: None, + partition_expr_version: None, + }), + ) + .await + .map(|_| ()) +} + +#[tokio::test] +async fn test_reopen_after_a_failed_create_shows_only_the_acknowledged_write() { + let mut env = TestEnv::with_prefix("object-store-wal-failed-create").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + let key_rows = |value| Rows { + schema: schema.clone(), + rows: build_rows_for_key("k", 0, 1, value), + }; + + // Object 1 is stored, but its create is reported as failed, so the write + // fails. + let failed = { + let engine = engine.clone(); + let rows = key_rows(1); + tokio::spawn(async move { try_put(&engine, REGION_A, rows).await }) + }; + tokio::time::timeout(WAIT, store.wait_for_admitted_appends(1)) + .await + .unwrap() + .unwrap(); + store.fail_next_create_after_write(); + store.seal_open_batch().await.unwrap_err(); + assert!(failed.await.unwrap().is_err()); + assert_eq!(vec![0, 1], wal_object_seqs(&object_store).await); + assert_eq!(0, latest(&store, REGION_A)); + + // The same key and timestamp is written again under the same row + // sequence, and acknowledged from object 2. + let mut writer = SealedWriter { + store: store.clone(), + admitted: 1, + }; + writer + .put_and_seal(&engine, vec![(REGION_A, key_rows(2))]) + .await; + // Object 2 carries the boundary past the start object, which is + // collected. + wait_for_collection(&store).await; + assert_eq!(vec![1, 2], wal_object_seqs(&object_store).await); + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + let region_a = region(&engine, REGION_A); + assert_eq!(1, region_a.version_control.committed_sequence()); + let expected = "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| k | 2.0 | 1970-01-01T00:00:00 | ++-------+---------+---------------------+"; + assert_eq!(expected, scan_rows(&engine, REGION_A).await); + + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // Object 1 is off the chain, so only the acknowledged write replays. + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 1, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(expected, scan_rows(&engine, REGION_A).await); +} + +/// Requests a flush of the region in the background and returns its handle; +/// the result is up to the caller. +fn spawn_flush( + engine: &MitoEngine, + region_id: RegionId, +) -> tokio::task::JoinHandle> { + let engine = engine.clone(); + tokio::spawn(async move { + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest { + row_group_size: None, + reason: None, + }), + ) + .await + .map(|_| ()) + }) +} + +/// Seals the open batch in the background; with creates held, the seal +/// returns only once they are released. +fn spawn_seal( + store: &Arc, +) -> tokio::task::JoinHandle> { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) +} + +/// Waits until a testing hook of the store observes the state it waits for. +async fn wait_for(hook: impl Future>) { + tokio::time::timeout(WAIT, hook) + .await + .expect("the store must reach the state") + .unwrap(); +} + +/// Tears the engine and its store down as a crash of the process would: the +/// store writes nothing more and fails every caller still waiting for it, +/// then the engine stops. Returns once nothing holds the store any more, so a +/// store reopened on the prefix never runs beside the old one. +async fn crash(engine: MitoEngine, store: Arc) { + store.crash().await; + engine.stop().await.unwrap(); + drop(engine); + let released = Arc::downgrade(&store); + drop(store); + tokio::time::timeout(WAIT, async { + while released.strong_count() > 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("nothing may hold the crashed store"); +} + +#[tokio::test] +async fn test_flush_waits_until_the_wal_is_durable() { + let mut env = TestEnv::with_prefix("object-store-wal-flush-barrier").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema) = create_region(&engine, REGION_A, &[]).await; + + // The write is acknowledged on admission; its object is not created + // while creates are held. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(1, 1), + memtable_rows: 3, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(0, latest(&store, REGION_A)); + + // The flush writes its SST, then waits for the entry to be durable + // before it records the entry as flushed in the manifest. + let flush = spawn_flush(&engine, REGION_A); + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!flush.is_finished()); + assert!(!seal.is_finished()); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_A) + .await + .manifest_flushed_entry_id + ); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + store.release_creates(); + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .unwrap(); + tokio::time::timeout(WAIT, flush) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(1, 1), + last_entry_id: entry_id(1, 1), + topic_latest_entry_id: entry_id(1, 1), + manifest_flushed_entry_id: entry_id(1, 1), + memtable_rows: 0, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(3, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_enqueued_crash_before_the_object_exists_replays_durable_entries() { + let mut env = TestEnv::with_prefix("object-store-wal-enqueued-crash").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is acknowledged and a flush is attempted, but its object is + // never created: the flush waits and the manifest keeps watermark 0. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + let flush = spawn_flush(&engine, REGION_A); + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!flush.is_finished()); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_A) + .await + .manifest_flushed_entry_id + ); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // The process dies with the create still held: the waiting flush and + // the seal fail, and the object is never written. + store.crash().await; + assert!(flush.await.unwrap().is_err()); + assert!(seal.await.unwrap().is_err()); + crash(engine, store).await; + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // Nothing is durable, so nothing replays. The restart writes start + // object 1 and the region's next entry lands in object 2; the lost entry + // was inside the unpersisted backlog. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + EMPTY_RECOVERY_STATE, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(0, latest(&store, REGION_A)); + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 3, 5))]) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!( + entry_id(2, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + // Opening the region recorded its watermark, so object 2 carries the + // boundary past both start objects. + wait_for_collection(&store).await; + assert_eq!(vec![2], wal_object_seqs(&object_store).await); + let rows_before = scan_rows(&engine, REGION_A).await; + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); + drop(writer); + crash(engine, store).await; + + // The second restart replays the entry that became durable and skips + // nothing: the manifest never named an entry that was not durable. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 2, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_drop_cancels_a_flush_waiting_for_wal_durability() { + let mut env = TestEnv::with_prefix("object-store-wal-flush-barrier-drop").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema) = create_region(&engine, REGION_A, &[]).await; + + // The flush reaches the durability barrier and waits for a create that + // is held. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + let flush = spawn_flush(&engine, REGION_A); + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!flush.is_finished()); + + // The drop cancels the flush instead of waiting behind the upload. + tokio::time::timeout( + WAIT, + engine.handle_request( + REGION_A, + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }), + ), + ) + .await + .unwrap() + .unwrap(); + assert!(!engine.is_region_exists(REGION_A)); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + assert!( + tokio::time::timeout(WAIT, flush) + .await + .unwrap() + .unwrap() + .is_err() + ); + + store.release_creates(); + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .unwrap(); +} + +#[rstest] +#[case(RegionTruncateRequest::All)] +#[case(RegionTruncateRequest::Unflushed)] +#[tokio::test] +async fn test_enqueued_truncate_waits_until_the_wal_is_durable( + #[case] request: RegionTruncateRequest, +) { + let mut env = TestEnv::with_prefix("object-store-wal-truncate-barrier").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is acknowledged and a truncate is attempted, but its object + // is never created: the truncate waits and the manifest keeps frontier 0. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + let truncate = { + let engine = engine.clone(); + tokio::spawn(async move { + engine + .handle_request(REGION_A, RegionRequest::Truncate(request)) + .await + .map(|_| ()) + }) + }; + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!truncate.is_finished()); + let manifest = region(&engine, REGION_A).manifest_ctx.manifest().await; + assert_eq!(0, manifest.flushed_entry_id); + assert_eq!(None, manifest.truncated_entry_id); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // The process dies with the create still held: the waiting truncate and + // the seal fail. + store.crash().await; + assert!(truncate.await.unwrap().is_err()); + assert!(seal.await.unwrap().is_err()); + crash(engine, store).await; + + // Nothing is durable and the manifest names no entry, so the region + // starts over after start object 1, and its next entry becomes durable + // in object 2. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + EMPTY_RECOVERY_STATE, + region_recovery_state(&engine, REGION_A).await + ); + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 3, 5))]) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + let rows_before = scan_rows(&engine, REGION_A).await; + drop(writer); + crash(engine, store).await; + + // The second restart replays the durable entry instead of skipping it. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 2, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[rstest] +#[case(RegionTruncateRequest::All)] +#[case(RegionTruncateRequest::Unflushed)] +#[tokio::test] +async fn test_enqueued_truncate_completes_once_the_wal_is_durable( + #[case] request: RegionTruncateRequest, +) { + let mut env = TestEnv::with_prefix("object-store-wal-truncate-durable").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + let truncated_entry_id = + matches!(request, RegionTruncateRequest::All).then_some(entry_id(1, 1)); + + // The truncate waits for the acknowledged entry, then records it once + // its object is created. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + let truncate = { + let engine = engine.clone(); + tokio::spawn(async move { + engine + .handle_request(REGION_A, RegionRequest::Truncate(request)) + .await + .map(|_| ()) + }) + }; + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!truncate.is_finished()); + store.release_creates(); + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .unwrap(); + tokio::time::timeout(WAIT, truncate) + .await + .unwrap() + .unwrap() + .unwrap(); + let manifest = region(&engine, REGION_A).manifest_ctx.manifest().await; + assert_eq!(entry_id(1, 1), manifest.flushed_entry_id); + assert_eq!(truncated_entry_id, manifest.truncated_entry_id); + assert_eq!(0, engine.get_region_statistic(REGION_A).unwrap().num_rows); + engine.stop().await.unwrap(); + drop(engine); + drop(store); + + // The restart replays nothing: the discarded rows stay discarded. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(1, 1), + last_entry_id: entry_id(1, 1), + topic_latest_entry_id: entry_id(1, 1), + manifest_flushed_entry_id: entry_id(1, 1), + memtable_rows: 0, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(0, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_flush_does_not_publish_a_frontier_for_a_lost_enqueued_backlog() { + let mut env = TestEnv::with_prefix("object-store-wal-lost-backlog").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is acknowledged, then its create fails while stop has begun + // but the stop command is not handled yet: the backlog is lost. + store.hold_creates(); + store.fail_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + store.begin_stop(); + store.release_creates(); + assert!( + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .is_err() + ); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // A flush in that window fails at the barrier instead of recording the + // lost entry as flushed. + let flush = spawn_flush(&engine, REGION_A); + assert!( + tokio::time::timeout(WAIT, flush) + .await + .unwrap() + .unwrap() + .is_err() + ); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_A) + .await + .manifest_flushed_entry_id + ); + assert!(store.stop().await.is_err()); +} diff --git a/src/mito2/src/engine/object_store_wal_test.rs b/src/mito2/src/engine/object_store_wal_test.rs new file mode 100644 index 00000000000..d4225b6fcf4 --- /dev/null +++ b/src/mito2/src/engine/object_store_wal_test.rs @@ -0,0 +1,379 @@ +// 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. + +//! Tests for regions whose WAL options select the object store WAL. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use api::v1::Rows; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_recordbatch::RecordBatches; +use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; +use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use futures::TryStreamExt; +use log_store::ObjectStoreLogStore; +use store_api::logstore::LogStore; +use store_api::logstore::provider::Provider; +use store_api::region_engine::RegionEngine; +use store_api::region_request::{ + PathType, RegionCleanUpRequest, RegionCloseRequest, RegionDropRequest, RegionOpenRequest, + RegionRequest, +}; +use store_api::storage::{RegionId, ScanRequest}; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, +}; +use crate::wal::EntryId; + +fn wal_config() -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + prefix: "cluster-a/wal".to_string(), + ..Default::default() + } +} + +/// The node prefix a standalone datanode derives from the configured root, +/// which the regions persist in their WAL options. +fn node_prefix() -> String { + wal_config().node_prefix(0, STANDALONE_GENERATION) +} + +fn wal_options() -> HashMap { + let options = WalOptions::ObjectStore(ObjectStoreWalOptions::new(node_prefix())); + HashMap::from([( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&options).unwrap(), + )]) +} + +fn object_store_provider(region_id: RegionId) -> Provider { + Provider::object_store_provider(region_id, node_prefix()) +} + +fn open_request(table_dir: String) -> RegionOpenRequest { + RegionOpenRequest { + engine: String::new(), + table_dir, + options: wal_options(), + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + } +} + +async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc) { + env.create_engine_with_object_store_wal(MitoConfig::default(), &wal_config()) + .await +} + +/// Creates a region with object store WAL options and returns its table dir and row schema. +async fn create_region( + engine: &MitoEngine, + region_id: RegionId, +) -> (String, Vec) { + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]) + .build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + (table_dir, schema) +} + +/// Creates a region and writes rows `[start, end)` to it in one WAL entry. +async fn create_region_with_rows( + engine: &MitoEngine, + region_id: RegionId, + start: usize, + end: usize, +) -> String { + let (table_dir, schema) = create_region(engine, region_id).await; + put_rows( + engine, + region_id, + Rows { + schema, + rows: build_rows(start, end), + }, + ) + .await; + table_dir +} + +async fn close_region(engine: &MitoEngine, region_id: RegionId) { + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); +} + +async fn count_rows(engine: &MitoEngine, region_id: RegionId) -> usize { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + RecordBatches::try_collect(stream) + .await + .unwrap() + .iter() + .map(|batch| batch.num_rows()) + .sum() +} + +/// Returns the ids of the entries the store still hands out for `region_id`. +async fn stored_entry_ids(log_store: &ObjectStoreLogStore, region_id: RegionId) -> Vec { + log_store + .read(&object_store_provider(region_id), 0, None) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .map(|entry| entry.entry_id()) + .collect() +} + +#[tokio::test] +async fn test_create_and_reopen_region_with_object_store_wal() { + let mut env = TestEnv::with_prefix("object-store-wal-reopen").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_id = RegionId::new(1, 1); + let provider = object_store_provider(region_id); + + let table_dir = create_region_with_rows(&engine, region_id, 0, 3).await; + assert_eq!(provider, engine.get_region(region_id).unwrap().provider); + let last_entry_id = log_store.latest_entry_id(&provider).unwrap(); + assert_eq!( + vec![last_entry_id], + stored_entry_ids(&log_store, region_id).await + ); + + reopen_region(&engine, region_id, table_dir, true, wal_options()).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!(provider, region.provider); + assert_eq!( + last_entry_id, + region.version_control.current().last_entry_id + ); + assert_eq!(3, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_topic_latest_entry_id_follows_replay() { + let mut env = TestEnv::with_prefix("object-store-wal-latest-entry-id").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_id = RegionId::new(1, 1); + let provider = object_store_provider(region_id); + + let (table_dir, schema) = create_region(&engine, region_id).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!(0, region.version_control.current().version.flushed_entry_id); + assert_eq!(0, region.topic_latest_entry_id.load(Ordering::Relaxed)); + + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let flushed_entry_id = region.version_control.current().version.flushed_entry_id; + assert_eq!( + flushed_entry_id, + log_store.latest_entry_id(&provider).unwrap() + ); + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows(3, 5), + }, + ) + .await; + let latest_entry_id = log_store.latest_entry_id(&provider).unwrap(); + assert!(latest_entry_id > flushed_entry_id); + + // Replay refills the memtable with the unflushed entry, so the topic latest + // entry id stays at the replay start. + reopen_region(&engine, region_id, table_dir.clone(), true, wal_options()).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!( + flushed_entry_id, + region.version_control.current().version.flushed_entry_id + ); + assert_eq!( + latest_entry_id, + region.version_control.current().last_entry_id + ); + assert_eq!( + flushed_entry_id, + region.topic_latest_entry_id.load(Ordering::Relaxed) + ); + assert_eq!(5, count_rows(&engine, region_id).await); + + // Nothing is replayed after a flush, so the topic latest entry id comes from the store. + flush_region(&engine, region_id, None).await; + reopen_region(&engine, region_id, table_dir, true, wal_options()).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!( + latest_entry_id, + region.version_control.current().version.flushed_entry_id + ); + assert_eq!( + latest_entry_id, + region.topic_latest_entry_id.load(Ordering::Relaxed) + ); + assert_eq!(5, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_open_region_replays_only_its_own_entries() { + let mut env = TestEnv::with_prefix("object-store-wal-isolation").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(1, 2); + + let table_dir_a = create_region_with_rows(&engine, region_a, 0, 3).await; + let table_dir_b = create_region_with_rows(&engine, region_b, 0, 5).await; + close_region(&engine, region_a).await; + close_region(&engine, region_b).await; + // The store holds the entries of both regions under the same prefix. + assert_eq!(1, stored_entry_ids(&log_store, region_a).await.len()); + assert_eq!(1, stored_entry_ids(&log_store, region_b).await.len()); + + engine + .handle_request(region_a, RegionRequest::Open(open_request(table_dir_a))) + .await + .unwrap(); + assert_eq!(3, count_rows(&engine, region_a).await); + assert!(!engine.is_region_exists(region_b)); + + engine + .handle_request(region_b, RegionRequest::Open(open_request(table_dir_b))) + .await + .unwrap(); + assert_eq!(5, count_rows(&engine, region_b).await); + assert_eq!(3, count_rows(&engine, region_a).await); +} + +#[tokio::test] +async fn test_batch_open_regions_with_object_store_wal() { + let mut env = TestEnv::with_prefix("object-store-wal-batch-open").await; + let (engine, _) = new_engine(&mut env).await; + + let mut regions = Vec::new(); + for i in 1..=3 { + let region_id = RegionId::new(1, i); + let num_rows = i as usize + 1; + let table_dir = create_region_with_rows(&engine, region_id, 0, num_rows).await; + close_region(&engine, region_id).await; + regions.push((region_id, table_dir, num_rows)); + } + + let requests = regions + .iter() + .map(|(region_id, table_dir, _)| (*region_id, open_request(table_dir.clone()))) + .collect(); + let responses = engine + .handle_batch_open_requests(4, requests) + .await + .unwrap(); + assert_eq!(regions.len(), responses.len()); + for (_, response) in responses { + response.unwrap(); + } + + for (region_id, _, num_rows) in regions { + assert_eq!( + object_store_provider(region_id), + engine.get_region(region_id).unwrap().provider + ); + assert_eq!(num_rows, count_rows(&engine, region_id).await); + } +} + +#[tokio::test] +async fn test_drop_and_offline_cleanup_with_object_store_wal() { + let mut env = TestEnv::with_prefix("object-store-wal-cleanup").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(1, 2); + + let _ = create_region_with_rows(&engine, region_a, 0, 3).await; + let table_dir_b = create_region_with_rows(&engine, region_b, 0, 3).await; + + engine + .handle_request( + region_a, + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }), + ) + .await + .unwrap(); + // The drop obsoletes the entries of region A only. + assert!(stored_entry_ids(&log_store, region_a).await.is_empty()); + assert_eq!(1, stored_entry_ids(&log_store, region_b).await.len()); + + close_region(&engine, region_b).await; + engine + .handle_request( + region_b, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: String::new(), + table_dir: table_dir_b, + path_type: PathType::Bare, + options: wal_options(), + }), + ) + .await + .unwrap(); + assert!(stored_entry_ids(&log_store, region_b).await.is_empty()); +} + +#[tokio::test] +async fn test_object_store_wal_options_reject_raft_engine_log_store() { + let mut env = TestEnv::with_prefix("object-store-wal-raft-engine").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]) + .build(); + let err = engine + .handle_request(RegionId::new(1, 1), RegionRequest::Create(request)) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); +} diff --git a/src/mito2/src/engine/prune_test.rs b/src/mito2/src/engine/prune_test.rs index f2171510345..71e6845e96f 100644 --- a/src/mito2/src/engine/prune_test.rs +++ b/src/mito2/src/engine/prune_test.rs @@ -288,7 +288,7 @@ async fn test_prune_memtable_complex_expr_with_format(flat_format: bool) { ) .await; - // ts filter will be ignored when pruning time series in memtable. + // Both the compound time range and the tag predicate are evaluated during the scan. let filters = vec![time_range_expr(4, 7), col("tag_0").lt(lit("6"))]; let stream = engine @@ -306,10 +306,6 @@ async fn test_prune_memtable_complex_expr_with_format(flat_format: bool) { +-------+---------+---------------------+ | tag_0 | field_0 | ts | +-------+---------+---------------------+ -| 0 | 0.0 | 1970-01-01T00:00:00 | -| 1 | 1.0 | 1970-01-01T00:00:01 | -| 2 | 2.0 | 1970-01-01T00:00:02 | -| 3 | 3.0 | 1970-01-01T00:00:03 | | 4 | 4.0 | 1970-01-01T00:00:04 | | 5 | 5.0 | 1970-01-01T00:00:05 | +-------+---------+---------------------+"; diff --git a/src/mito2/src/engine/puffin_index.rs b/src/mito2/src/engine/puffin_index.rs index 5624405dd7e..278d9017e7f 100644 --- a/src/mito2/src/engine/puffin_index.rs +++ b/src/mito2/src/engine/puffin_index.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::convert::TryFrom; +use std::sync::Arc; use common_base::range_read::RangeReader; use common_telemetry::warn; @@ -111,7 +112,7 @@ pub(crate) async fn collect_index_entries_from_puffin( ) .await; - let bloom_value = bloom_meta.as_ref().map(bloom_meta_value); + let bloom_value = bloom_meta.as_deref().map(bloom_meta_value); let (target_type, target_json) = decode_target_info(target_key); let meta_json = build_meta_json(bloom_value, None, None); let entry = build_index_entry( @@ -137,7 +138,7 @@ pub(crate) async fn collect_index_entries_from_puffin( ) .await; - let bloom_value = bloom_meta.as_ref().map(bloom_meta_value); + let bloom_value = bloom_meta.as_deref().map(bloom_meta_value); let fulltext_value = Some(fulltext_meta_value(blob)); let (target_type, target_json) = decode_target_info(target_key); let meta_json = build_meta_json(bloom_value, fulltext_value, None); @@ -294,7 +295,7 @@ async fn try_read_bloom_meta( cache: Option<&BloomFilterIndexCacheRef>, tag: Tag, context: &IndexEntryContext<'_>, -) -> Option { +) -> Option> { let column_id = decode_column_id(target_key); // Failures are logged but do not abort the overall metadata collection. diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index bf089737f35..f86b327735d 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -56,7 +56,7 @@ use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner}; use store_api::region_request::{ AlterKind, RegionAlterRequest, RegionBulkInsertsRequest, RegionCompactRequest, - RegionPutRequest, RegionRequest, SetRegionOption, + RegionOpenRequest, RegionPutRequest, RegionRequest, SetRegionOption, }; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ @@ -73,6 +73,215 @@ use crate::test_util; use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; use crate::test_util::{CreateRequestBuilder, TestEnv, reopen_region}; +#[tokio::test] +async fn test_nullable_label_filters_in_memtables_and_ssts() { + use datafusion_expr::{BinaryExpr as LogicalBinaryExpr, Cast, Expr}; + + for (flat, memtable) in [(false, "time_series"), (true, "bulk")] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", memtable) + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows = [None, Some(""), Some("tmpfs"), Some("ext4")] + .into_iter() + .enumerate() + .map(|(i, label)| { + let mut row = row(vec![ + ValueData::StringValue(String::new()), + ValueData::F64Value(i as f64), + ValueData::TimestampMillisecondValue(i as i64), + ]); + row.values[0].value_data = label.map(|v| ValueData::StringValue(v.into())); + row + }) + .collect(); + test_util::put_rows(&engine, region_id, Rows { schema, rows }).await; + for flushed in [false, true] { + if flushed { + test_util::flush_region(&engine, region_id, None).await; + } + let mut cases = Vec::new(); + for (op, literal, include_null, expected) in [ + (Operator::NotEq, "tmpfs", true, vec![0, 1, 3]), + ( + Operator::RegexNotMatch, + "^(?:tmpfs|vfat)$", + true, + vec![0, 1, 3], + ), + (Operator::Eq, "", true, vec![0, 1]), + (Operator::NotEq, "", false, vec![2, 3]), + (Operator::RegexMatch, "^(?:ext4)$", false, vec![3]), + ] { + for cast in [false, true] { + let column = if cast { + Expr::Cast(Cast::new(Box::new(col("tag_0")), DataType::Utf8)) + } else { + col("tag_0") + }; + let comparison = Expr::BinaryExpr(LogicalBinaryExpr::new( + Box::new(column), + op, + Box::new(lit(literal)), + )); + let filter = if include_null { + col("tag_0").is_null().or(comparison) + } else { + comparison + }; + cases.push((filter, expected.clone())); + } + } + let column = col("tag_0"); + cases.extend([ + (column.clone().is_null(), vec![0]), + (column.clone().is_not_null(), vec![1, 2, 3]), + ( + column + .clone() + .is_null() + .or(column.clone().not_eq(lit("tmpfs"))) + .and(column.clone().is_null().or(column.clone().not_eq(lit("")))), + vec![0, 3], + ), + ( + column.clone().eq(lit("ext4")).or(column.eq(lit("tmpfs"))), + vec![2, 3], + ), + ]); + for (filter, expected) in cases { + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + filters: vec![filter.clone()], + distribution: flat.then_some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut actual = Vec::new(); + for batch in batches.iter() { + let ts = batch + .column_by_name("ts") + .unwrap() + .as_primitive::(); + actual.extend(ts.values().iter().copied()); + } + actual.sort_unstable(); + assert_eq!(actual, expected, "flat={flat}, flushed={flushed}, {filter}"); + } + } + } +} + +#[tokio::test] +async fn test_compound_time_filters_before_last_row() { + for (flat, memtable) in [ + (false, "time_series"), + (true, "time_series"), + (true, "bulk"), + ] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", memtable) + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows = (0..5) + .map(|i| { + row(vec![ + ValueData::StringValue("series".into()), + ValueData::F64Value(i as f64), + ValueData::TimestampMillisecondValue(i), + ]) + }) + .collect(); + test_util::put_rows(&engine, region_id, Rows { schema, rows }).await; + let ts = |value| lit(ScalarValue::TimestampMillisecond(Some(value), None)); + let cases = [ + (col("ts").eq(ts(1)).or(col("ts").eq(ts(3))), vec![1, 3]), + (col("ts").gt_eq(ts(1)).and(col("ts").lt(ts(3))), vec![1, 2]), + (col("ts").is_null().or(col("ts").lt(ts(2))), vec![0, 1]), + ]; + for flushed in [false, true] { + if flushed { + test_util::flush_region(&engine, region_id, None).await; + } + for (filter, expected) in &cases { + for selector in [ + None, + Some(TimeSeriesRowSelector::LastRow { after_merge: false }), + Some(TimeSeriesRowSelector::LastRow { after_merge: true }), + ] { + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + filters: vec![ + col("tag_0").is_null().or(col("tag_0").eq(lit("series"))), + filter.clone(), + ], + distribution: flat.then_some(TimeSeriesDistribution::PerSeries), + series_row_selector: selector, + ..Default::default() + }, + ) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut actual = Vec::new(); + for batch in batches.iter() { + actual.extend( + batch + .column_by_name("ts") + .unwrap() + .as_primitive::() + .values() + .iter() + .copied(), + ); + } + actual.sort_unstable(); + let expected = if selector.is_some() { + expected.last().copied().into_iter().collect::>() + } else { + expected.clone() + }; + assert_eq!( + actual, expected, + "flat={flat}, memtable={memtable}, flushed={flushed}, selector={selector:?}, {filter}" + ); + } + } + } + } +} + #[tokio::test] async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResult<()> { // Create a region with a JSON2 field whose physical Parquet representation is a nested struct. @@ -1340,6 +1549,7 @@ async fn check_two_phase_series_scan( let store = region.series_index_store.clone().unwrap(); let sequence = region.flushed_sequence(); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_millisecond(0), bucket_end: Timestamp::new_millisecond(2001), @@ -1412,7 +1622,13 @@ async fn check_two_phase_series_scan( .unwrap(); writer.write(0, &batch).await.unwrap(); writer.finish().await.unwrap(); - index_version.range_indexes.insert(file_id); + index_version.range_indexes.insert( + file_id, + crate::series_index::RangeIndexEntry { + file_id, + file_size: 0, + }, + ); } region .series_index_version_control @@ -1607,7 +1823,7 @@ async fn check_two_phase_series_scan( // succeed even when that file is unavailable. let region = engine.find_region(region_id).unwrap(); let version = region.series_index_version_control.current(); - let file_id = *version.range_indexes.iter().next().unwrap(); + let file_id = *version.range_indexes.keys().next().unwrap(); region .series_index_store .as_ref() @@ -2775,13 +2991,21 @@ 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; + check_bulk_skip_wal_recovery(false, false, false).await; + check_bulk_skip_wal_recovery(false, true, false).await; + check_bulk_skip_wal_recovery(true, false, false).await; + check_bulk_skip_wal_recovery(true, true, false).await; } -async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { +#[tokio::test] +async fn test_bulk_skip_wal_close_recovery() { + check_bulk_skip_wal_recovery(false, false, true).await; + check_bulk_skip_wal_recovery(false, true, true).await; + check_bulk_skip_wal_recovery(true, false, true).await; + check_bulk_skip_wal_recovery(true, true, true).await; +} + +async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool, close: bool) { let mut env = TestEnv::new().await; let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -2832,7 +3056,41 @@ async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { u64::from(!skip_wal) ); } - reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + let engine = if close { + reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + assert_eq!( + region.version().flushed_sequence, + if flush || skip_wal { 4 } else { 0 } + ); + engine + } else { + // Stopping the engine directly exercises recovery without close-time flush. + let engine = env.reopen_engine(engine, MitoConfig::default()).await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: store_api::region_request::PathType::Bare, + options: HashMap::new(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + }; + assert!( + !engine + .get_region(region_id) + .unwrap() + .version() + .memtables + .has_skip_wal_writes() + ); let stream = engine .scan_to_stream(region_id, ScanRequest::default()) .await @@ -2844,7 +3102,7 @@ async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { .iter() .map(|b| b.num_rows()) .sum::(), - if skip_wal && !flush { 0 } else { 4 } + if skip_wal && !flush && !close { 0 } else { 4 } ); } diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index 363318d1cf5..915099e7228 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -34,7 +35,7 @@ use crate::engine::listener::AlterFlushListener; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, TestEnv, build_rows, flush_region, kafka_log_store_factory, multiple_log_store_factories, prepare_test_for_kafka_log_store, - put_rows, raft_engine_log_store_factory, rows_schema, + put_rows, raft_engine_log_store_factory, reopen_region, rows_schema, }; fn set_skip_wal_request(skip_wal: bool) -> RegionRequest { @@ -920,3 +921,192 @@ async fn test_close_region_after_truncate_skip_wal() { let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(3, total_rows); } + +#[rstest] +#[case(false)] +#[case(true)] +#[tokio::test] +async fn test_close_request_skip_wal( + #[case] skip_wal: bool, + #[values(false, true)] flat_format: bool, +) { + let mut env = TestEnv::with_prefix("close-request-skip-wal").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(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + assert!(!region.skip_wal()); + assert_eq!(region.version().memtables.has_skip_wal_writes(), skip_wal); + let expected = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + + reopen_region(&engine, region_id, request.table_dir, false, HashMap::new()).await; + // Logged writes are recovered without a close-time flush. + assert_eq!( + region.version().flushed_sequence, + if skip_wal { 3 } else { 0 } + ); + let reopened = engine.get_region(region_id).unwrap(); + assert!(!reopened.version().memtables.has_skip_wal_writes()); + let actual = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(actual, expected); +} + +#[rstest] +#[case(false)] +#[case(true)] +#[tokio::test] +async fn test_close_request_skip_wal_during_flush(#[case] new_skip_wal: bool) { + let mut env = TestEnv::with_prefix("close-request-skip-wal-during-flush").await; + let listener = Arc::new(AlterFlushListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: true, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let engine_cloned = engine.clone(); + let flush_job = tokio::spawn(async move { + flush_region(&engine_cloned, region_id, None).await; + }); + tokio::time::timeout(Duration::from_secs(5), listener.wait_flush_begin()) + .await + .unwrap(); + + // These rows belong to the next mutable, while the skipped rows are frozen. + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: new_skip_wal, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(3, 6), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let expected = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + let request_count = listener.request_count(); + let engine_cloned = engine.clone(); + let close_job = tokio::spawn(async move { + engine_cloned + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + }); + tokio::time::timeout( + Duration::from_secs(5), + listener.wait_request_count(request_count + 1), + ) + .await + .unwrap(); + listener.wake_flush(); + tokio::time::timeout(Duration::from_secs(5), flush_job) + .await + .unwrap() + .unwrap(); + if new_skip_wal { + tokio::time::timeout(Duration::from_secs(5), listener.wait_flush_begin()) + .await + .unwrap(); + listener.wake_flush(); + } + tokio::time::timeout(Duration::from_secs(5), close_job) + .await + .unwrap() + .unwrap(); + assert_eq!( + region.version().flushed_sequence, + if new_skip_wal { 6 } else { 3 } + ); + assert!(!region.version().memtables.has_skip_wal_writes()); + + reopen_region(&engine, region_id, request.table_dir, false, HashMap::new()).await; + let actual = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(actual, expected); +} diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 1592bfebd60..f738044c9ef 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -453,6 +453,14 @@ pub enum Error { source: BoxedError, }, + #[snafu(display("Failed to wait for the WAL to be durable, region_id: {}", region_id))] + WaitWalDurable { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + source: BoxedError, + }, + // Shared error for each writer in the write group. #[snafu(display("Failed to write region"))] WriteGroup { source: Arc }, @@ -773,14 +781,6 @@ pub enum Error { location: Location, }, - #[cfg(feature = "vector_index")] - #[snafu(display("Failed to apply vector index: {}", reason))] - ApplyVectorIndex { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Failed to push index value"))] PushIndexValue { source: index::inverted_index::error::Error, @@ -1197,22 +1197,6 @@ pub enum Error { location: Location, }, - #[cfg(feature = "vector_index")] - #[snafu(display("Failed to build vector index: {}", reason))] - VectorIndexBuild { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[cfg(feature = "vector_index")] - #[snafu(display("Failed to finish vector index: {}", reason))] - VectorIndexFinish { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Manual compaction is override by following operations."))] ManualCompactionOverride {}, @@ -1463,9 +1447,10 @@ impl ErrorExt for Error { OpenDal { .. } | ManifestDeltaNotFound { .. } | ReadParquet { .. } => { StatusCode::StorageUnavailable } - WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } => { - source.status_code() - } + WriteWal { source, .. } + | ReadWal { source, .. } + | DeleteWal { source, .. } + | WaitWalDurable { source, .. } => source.status_code(), CompressObject { .. } | DecompressObject { .. } | SerdeJson { .. } @@ -1571,8 +1556,6 @@ impl ErrorExt for Error { | PushIndexValue { source, .. } | ApplyInvertedIndex { source, .. } | IndexFinish { source, .. } => source.status_code(), - #[cfg(feature = "vector_index")] - ApplyVectorIndex { .. } => StatusCode::Internal, PuffinReadBlob { source, .. } | PuffinAddBlob { source, .. } | PuffinInitStager { source, .. } @@ -1615,9 +1598,6 @@ impl ErrorExt for Error { source.status_code() } - #[cfg(feature = "vector_index")] - VectorIndexBuild { .. } | VectorIndexFinish { .. } => StatusCode::Internal, - ManualCompactionOverride {} | CompactionCancelled {} | FlushCancelled {} => { StatusCode::Cancelled } @@ -1680,6 +1660,7 @@ impl ErrorExt for Error { WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } + | WaitWalDurable { source, .. } | FetchManifests { source, .. } | External { source, .. } => source.retry_hint(), diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 8a96276e5f7..e910c81f896 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -69,6 +69,7 @@ use crate::sst::parquet::{ DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions, flat_format, }; use crate::sst::{FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema}; +use crate::wal::DurabilityBarrier; use crate::worker::WorkerListener; /// Global write buffer (memtable) manager. @@ -231,7 +232,7 @@ pub enum FlushReason { Repartition, /// Flush triggered by remote WAL pruning. RemoteWalPrune, - /// Flush before closing a Noop WAL region. + /// Flush before closing a region. Closing, } @@ -282,6 +283,8 @@ pub(crate) struct RegionFlushTask { /// /// This is used to generate the file meta. pub(crate) partition_expr: Option, + /// Waits until the WAL is durable through the entry id the flush records. + pub(crate) durability_barrier: DurabilityBarrier, } struct FlushTaskWaiters { @@ -509,6 +512,17 @@ impl RegionFlushTask { .await; } + // The manifest may name only a durable entry as flushed. A log store + // that acknowledges appends before their entries are durable can have + // handed out `last_entry_id` for an entry that is still in its backlog. + // A DDL that cancels the flush must not wait behind that upload. + let durable = self.durability_barrier.wait(version_data.last_entry_id); + tokio::pin!(durable); + match CancellableFuture::new(durable.as_mut(), state.cancel_handle()).await { + Ok(result) => result?, + Err(_) => return FlushCancelledSnafu.fail(), + } + let edit = RegionEdit { files_to_add: file_metas, files_to_remove: Vec::new(), @@ -836,7 +850,6 @@ impl RegionFlushTask { metadata: version.metadata.clone(), source, cache_manager: self.cache_manager.clone(), - storage: version.options.storage.clone(), max_sequence: Some(max_sequence), sst_write_format: if flat_format { FormatType::Flat @@ -849,8 +862,6 @@ impl RegionFlushTask { inverted_index_config: self.engine_config.inverted_index.clone(), fulltext_index_config: self.engine_config.fulltext_index.clone(), bloom_filter_index_config: self.engine_config.bloom_filter_index.clone(), - #[cfg(feature = "vector_index")] - vector_index_config: self.engine_config.vector_index.clone(), } } @@ -1706,6 +1717,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), } } @@ -1856,6 +1868,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), }; task.push_sender(OptionOutputTx::from(output_tx)); scheduler @@ -2146,6 +2159,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), }) .collect(); // Schedule first task. @@ -2450,6 +2464,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), }) .collect(); // Schedule first task. diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index 7bcf511f1a0..2a66294f59b 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -518,13 +518,9 @@ async fn write_and_list_entry(store: &ObjectStore, path: &str) -> Entry { .await .unwrap(); // List the parent directory to get the entry. - let parent = std::path::Path::new(path).parent().and_then(|p| { - if p.as_os_str().is_empty() { - None - } else { - Some(p) - } - }); + let parent = std::path::Path::new(path) + .parent() + .filter(|p| !p.as_os_str().is_empty()); let prefix = match parent { Some(p) => format!("{}/", p.to_str().unwrap()), None => String::new(), // root diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index 4f4409afd8a..dcbfe60dea5 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -16,10 +16,6 @@ //! //! Mito is the a region engine to store timeseries data. -#![feature(debug_closure_helpers)] -#![feature(duration_constructors)] -#![feature(binary_heap_pop_if)] - #[cfg(any(test, feature = "test"))] #[cfg_attr(feature = "test", allow(unused))] pub mod test_util; diff --git a/src/mito2/src/manifest/manager.rs b/src/mito2/src/manifest/manager.rs index c6a0d3f3228..436488f3c3c 100644 --- a/src/mito2/src/manifest/manager.rs +++ b/src/mito2/src/manifest/manager.rs @@ -49,7 +49,7 @@ pub struct RegionManifestOptions { pub manifest_dir: String, pub object_store: ObjectStore, pub compress_type: CompressionType, - /// Interval of version ([ManifestVersion](store_api::manifest::ManifestVersion)) between two checkpoints. + /// Interval of version ([ManifestVersion](store_api::ManifestVersion)) between two checkpoints. /// Set to 0 to disable checkpoint. pub checkpoint_distance: u64, pub remove_file_options: RemoveFileOptions, diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 94ccac472df..afede8904d8 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1872,7 +1872,7 @@ mod tests { let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum(); assert_eq!(5, total_rows); - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.num_rows() > 0); assert!(range.is_record_batch()); @@ -2468,7 +2468,7 @@ mod tests { let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum(); assert_eq!(10, total_rows); - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.num_rows() > 0); assert!(range.is_record_batch()); @@ -2554,7 +2554,7 @@ mod tests { // Read all data and verify let mut total_rows_read = 0; - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.is_record_batch()); let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); @@ -2642,7 +2642,7 @@ mod tests { assert_eq!(13, total_rows); let mut total_rows_read = 0; - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); for batch_result in record_batch_iter { let batch = batch_result.unwrap(); @@ -3063,7 +3063,7 @@ mod tests { // Read all data let mut total_rows_read = 0; - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.is_record_batch()); let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); @@ -3120,7 +3120,7 @@ mod tests { // Should return ranges but each range should produce an empty iterator // instead of an error. - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.is_record_batch()); let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); let total_rows: usize = record_batch_iter.map(|r| r.unwrap().num_rows()).sum(); diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index f994f242273..444cc59e317 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -20,7 +20,7 @@ use std::time::{Duration, Instant}; use api::helper::{ColumnDataTypeWrapper, to_grpc_value}; use api::v1::bulk_wal_entry::Body; -use api::v1::{ArrowIpc, BulkWalEntry, Mutation, OpType}; +use api::v1::{ArrowIpc, BulkWalEntry, Mutation, OpType, SemanticType}; use bytes::Bytes; use common_grpc::flight::{FlightDecoder, FlightEncoder, FlightMessage}; use common_recordbatch::DfRecordBatch as RecordBatch; @@ -71,7 +71,9 @@ 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, apply_float_field_encoding}; +use crate::sst::parquet::{ + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding, +}; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -218,6 +220,12 @@ impl BulkPart { // Finds columns that need to be filled let mut columns_to_fill = Vec::new(); for column_meta in ®ion_metadata.column_metadatas { + // Sparse bulk batches already carry tags in the encoded primary key. + if region_metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse + && column_meta.semantic_type == SemanticType::Tag + { + continue; + } // TODO(yingwen): Returns error if it is impure default after we support filling // bulk insert request in the frontend if !batch_columns.contains(column_meta.column_schema.name.as_str()) { @@ -1328,7 +1336,7 @@ impl BulkPartEncoder { .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_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); props = apply_float_field_encoding(props, &metadata, float_field_encoding); let writer_props = Some(props.build()); diff --git a/src/mito2/src/memtable/time_partition.rs b/src/mito2/src/memtable/time_partition.rs index a6b27cdd0f3..0ca88ca9a70 100644 --- a/src/mito2/src/memtable/time_partition.rs +++ b/src/mito2/src/memtable/time_partition.rs @@ -43,7 +43,7 @@ use crate::memtable::{KeyValues, MemtableBuilderRef, MemtableId, MemtableRef}; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; /// Initial time window if not specified. -const INITIAL_TIME_WINDOW: Duration = Duration::from_days(1); +const INITIAL_TIME_WINDOW: Duration = Duration::from_secs(86400); /// A partition holds rows with timestamps between `[min, max)`. #[derive(Debug, Clone)] diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index 4e720df88df..fad746927d4 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -641,7 +641,13 @@ impl Iter { predicate .exprs() .iter() - .filter_map(SimpleFilterEvaluator::try_new) + .filter_map(|expr| { + SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| { + metadata + .column_by_name(name) + .map(|column| column.column_schema.data_type.as_arrow_type()) + }) + }) .collect::>() }) .unwrap_or_default(); diff --git a/src/mito2/src/memtable/version.rs b/src/mito2/src/memtable/version.rs index d707128c93a..cd5dae10315 100644 --- a/src/mito2/src/memtable/version.rs +++ b/src/mito2/src/memtable/version.rs @@ -15,6 +15,7 @@ //! Memtable version. use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use common_time::Timestamp; @@ -23,8 +24,9 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::SequenceNumber; use crate::error::Result; -use crate::memtable::time_partition::TimePartitionsRef; -use crate::memtable::{MemtableId, MemtableRef}; +use crate::memtable::bulk::part::BulkPart; +use crate::memtable::time_partition::TimePartitions; +use crate::memtable::{KeyValues, MemtableBuilderRef, MemtableId, MemtableRef}; pub(crate) type SmallMemtableVec = SmallVec<[MemtableRef; 2]>; @@ -32,26 +34,35 @@ pub(crate) type SmallMemtableVec = SmallVec<[MemtableRef; 2]>; #[derive(Debug, Clone)] pub(crate) struct MemtableVersion { /// Mutable memtable. - pub(crate) mutable: TimePartitionsRef, + pub(crate) mutable: MutableMemtablesRef, /// Immutable memtables. /// /// We only allow one flush job per region but if a flush job failed, then we /// might need to store more than one immutable memtable on the next time we /// flush the region. immutables: SmallMemtableVec, + /// Whether immutable memtables may contain skip-WAL writes. + /// Flushes are serialized per region and cover all immutable memtables. + immutable_has_skip_wal_writes: bool, } pub(crate) type MemtableVersionRef = Arc; impl MemtableVersion { /// Returns a new [MemtableVersion] with specific mutable memtable. - pub(crate) fn new(mutable: TimePartitionsRef) -> MemtableVersion { + pub(crate) fn new(mutable: TimePartitions) -> MemtableVersion { MemtableVersion { - mutable, + mutable: Arc::new(MutableMemtables::new(mutable)), immutables: SmallVec::new(), + immutable_has_skip_wal_writes: false, } } + /// Returns whether current memtables may contain skip-WAL writes. + pub(crate) fn has_skip_wal_writes(&self) -> bool { + self.mutable.has_skip_wal_writes() || self.immutable_has_skip_wal_writes + } + /// Immutable memtables. pub(crate) fn immutables(&self) -> &[MemtableRef] { &self.immutables @@ -94,21 +105,25 @@ impl MemtableVersion { } // Update the time window. - let mutable = self.mutable.new_with_part_duration(time_window, None); + let mutable = self + .mutable + .partitions + .new_with_part_duration(time_window, None); common_telemetry::debug!( "Freeze empty memtable, update partition duration from {:?} to {:?}", self.mutable.part_duration(), time_window ); return Ok(Some(MemtableVersion { - mutable: Arc::new(mutable), + mutable: Arc::new(MutableMemtables::new(mutable)), immutables: self.immutables.clone(), + immutable_has_skip_wal_writes: self.immutable_has_skip_wal_writes, })); } // Marks the mutable memtable as immutable so it can free the memory usage from our // soft limit. - self.mutable.freeze()?; + self.mutable.partitions.freeze()?; // Fork the memtable. if Some(self.mutable.part_duration()) != time_window { common_telemetry::debug!( @@ -117,17 +132,23 @@ impl MemtableVersion { time_window ); } - let mutable = Arc::new(self.mutable.fork(metadata, time_window)); + let mutable = Arc::new(MutableMemtables::new( + self.mutable.partitions.fork(metadata, time_window), + )); let mut immutables = SmallVec::with_capacity(self.immutables.len() + self.mutable.num_partitions()); immutables.extend(self.immutables.iter().cloned()); // Pushes the mutable memtable to immutable list. - self.mutable.list_memtables_to_small_vec(&mut immutables); + self.mutable + .partitions + .list_memtables_to_small_vec(&mut immutables); Ok(Some(MemtableVersion { mutable, immutables, + immutable_has_skip_wal_writes: self.immutable_has_skip_wal_writes + || self.mutable.has_skip_wal_writes(), })) } @@ -139,6 +160,9 @@ impl MemtableVersion { .filter(|mem| !ids.contains(&mem.id())) .cloned() .collect(); + if self.immutables.is_empty() { + self.immutable_has_skip_wal_writes = false; + } } /// Returns the memory usage of the mutable memtable. @@ -182,3 +206,213 @@ impl MemtableVersion { self.mutable.is_empty() && self.immutables.is_empty() } } + +/// Mutable time partitions and their shared write state. +/// +/// Version clones share this object. Freezing creates a new object so writes to +/// the new mutable memtables cannot change the old version's state. +#[derive(Debug)] +pub(crate) struct MutableMemtables { + partitions: TimePartitions, + /// Set before installation because failed writes may install some rows. + has_skip_wal_writes: AtomicBool, +} + +pub(crate) type MutableMemtablesRef = Arc; + +impl MutableMemtables { + fn new(partitions: TimePartitions) -> Self { + Self { + partitions, + has_skip_wal_writes: AtomicBool::new(false), + } + } + + /// Marks that these memtables may contain skip-WAL writes. + pub(crate) fn mark_skip_wal_writes(&self) { + self.has_skip_wal_writes.store(true, Ordering::Relaxed); + } + + fn has_skip_wal_writes(&self) -> bool { + self.has_skip_wal_writes.load(Ordering::Relaxed) + } + + /// Writes key values to the mutable time partitions. + pub(crate) fn write(&self, kvs: &KeyValues) -> Result<()> { + self.partitions.write(kvs) + } + + /// Writes a bulk part to the mutable time partitions. + pub(crate) fn write_bulk(&self, part: BulkPart) -> Result<()> { + self.partitions.write_bulk(part) + } + + /// Returns whether all mutable time partitions are empty. + pub(crate) fn is_empty(&self) -> bool { + self.partitions.is_empty() + } + + /// Returns the memtable builder. + pub(crate) fn memtable_builder(&self) -> &MemtableBuilderRef { + self.partitions.memtable_builder() + } + + /// Returns the next memtable ID. + pub(crate) fn next_memtable_id(&self) -> MemtableId { + self.partitions.next_memtable_id() + } + + /// Returns the time partition duration. + pub(crate) fn part_duration(&self) -> Duration { + self.partitions.part_duration() + } + + /// Creates empty time partitions with an optional new duration and builder. + /// The returned partitions do not share data or write state with this object. + pub(crate) fn new_with_part_duration( + &self, + part_duration: Option, + memtable_builder: Option, + ) -> TimePartitions { + self.partitions + .new_with_part_duration(part_duration, memtable_builder) + } + + /// Returns the mutable memory usage. + pub(crate) fn memory_usage(&self) -> usize { + self.partitions.memory_usage() + } + + /// Returns the number of time series in mutable time partitions. + pub(crate) fn series_count(&self) -> usize { + self.partitions.series_count() + } + + fn num_rows(&self) -> u64 { + self.partitions.num_rows() + } + + fn num_partitions(&self) -> usize { + self.partitions.num_partitions() + } + + fn list_memtables(&self, memtables: &mut Vec) { + self.partitions.list_memtables(memtables); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memtable::time_series::TimeSeriesMemtableBuilder; + use crate::test_util::memtable_util; + + #[test] + fn test_skip_wal_writes_follow_memtable_lifecycle() { + for new_skip_wal_writes in [false, true] { + let metadata = memtable_util::metadata_for_test(); + let partitions = TimePartitions::new( + metadata.clone(), + Arc::new(TimeSeriesMemtableBuilder::default()), + 0, + Some(Duration::from_secs(5)), + ); + let version = MemtableVersion::new(partitions); + let snapshot = version.clone(); + assert!(!version.has_skip_wal_writes()); + + let kvs = memtable_util::build_key_values( + &metadata, + "hello".to_string(), + 0, + &[1000, 7000], + 1, + ); + version.mutable.mark_skip_wal_writes(); + version.mutable.write(&kvs).unwrap(); + assert!(snapshot.has_skip_wal_writes()); + + let mut frozen = version + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + assert_eq!(2, frozen.immutables().len()); + assert!(frozen.has_skip_wal_writes()); + assert!(!frozen.mutable.has_skip_wal_writes()); + assert!(!Arc::ptr_eq(&version.mutable, &frozen.mutable)); + + let during_flush = frozen.clone(); + if new_skip_wal_writes { + frozen.mutable.mark_skip_wal_writes(); + } + frozen.mutable.write(&kvs).unwrap(); + let ids: Vec<_> = frozen.immutables().iter().map(|mem| mem.id()).collect(); + frozen.remove_memtables(&ids); + assert_eq!(new_skip_wal_writes, frozen.has_skip_wal_writes()); + assert_eq!( + new_skip_wal_writes, + during_flush.mutable.has_skip_wal_writes() + ); + // Removing flushed memtables does not mutate the old flush snapshot. + assert_eq!(2, during_flush.immutables().len()); + assert!(during_flush.has_skip_wal_writes()); + + let mut next_flush = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + let ids: Vec<_> = next_flush.immutables().iter().map(|mem| mem.id()).collect(); + next_flush.remove_memtables(&ids); + assert!(!next_flush.has_skip_wal_writes()); + } + } + + #[test] + fn test_skip_wal_writes_survive_flush_retry() { + let metadata = memtable_util::metadata_for_test(); + let partitions = TimePartitions::new( + metadata.clone(), + Arc::new(TimeSeriesMemtableBuilder::default()), + 0, + Some(Duration::from_secs(5)), + ); + let version = MemtableVersion::new(partitions); + let kvs = memtable_util::build_key_values(&metadata, "hello".to_string(), 0, &[1000], 1); + version.mutable.mark_skip_wal_writes(); + version.mutable.write(&kvs).unwrap(); + let frozen = version + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + + // A failed flush leaves immutable memtables installed. Changing the + // empty mutable's partition duration must retain their write state. + let frozen = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(10))) + .unwrap() + .unwrap(); + assert!(frozen.has_skip_wal_writes()); + assert!(!frozen.mutable.has_skip_wal_writes()); + + // A retry includes the old immutable memtables and new logged writes. + frozen.mutable.write(&kvs).unwrap(); + let mut retried = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(10))) + .unwrap() + .unwrap(); + assert_eq!(2, retried.immutables().len()); + assert!(retried.has_skip_wal_writes()); + let ids: Vec<_> = retried.immutables().iter().map(|mem| mem.id()).collect(); + retried.remove_memtables(&ids); + assert!(!retried.has_skip_wal_writes()); + + let replacement = MemtableVersion::new(TimePartitions::new( + metadata, + frozen.mutable.memtable_builder().clone(), + frozen.mutable.next_memtable_id(), + None, + )); + assert!(!replacement.has_skip_wal_writes()); + assert!(frozen.has_skip_wal_writes()); + } +} diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index 49e91d93efd..4c275c248c9 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -318,6 +318,14 @@ lazy_static! { // Index metrics. lazy_static! { + /// Approximate published index bytes in open regions, refreshed by maintenance. + pub static ref SERIES_INDEX_DISK_BYTES: IntGauge = register_int_gauge!( + "greptime_mito_series_index_disk_bytes", "estimated series and range index bytes in open regions" + ).unwrap(); + /// Maintenance passes deferred by the estimated disk usage. + pub static ref SERIES_INDEX_CAPACITY_DEFERRED: IntCounter = register_int_counter!( + "greptime_mito_series_index_capacity_deferred_total", "series-index capacity deferrals" + ).unwrap(); // Index metrics. /// Outcomes of series-index reconciliation passes. pub static ref SERIES_INDEX_RECONCILE_TOTAL: IntCounterVec = register_int_counter_vec!( diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index 72e778851ee..65d356d79f2 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -45,7 +45,6 @@ use std::time::Duration; use api::v1::OpType; use arrow_schema::SchemaRef; -use async_trait::async_trait; use common_time::Timestamp; use datafusion_common::arrow::array::UInt8Array; use datatypes::arrow; @@ -65,7 +64,7 @@ use datatypes::vectors::{ }; use futures::TryStreamExt; use futures::stream::BoxStream; -use mito_codec::row_converter::{CompositeValues, PrimaryKeyCodec}; +use mito_codec::row_converter::{CompositeValues, DensePrimaryKeyCodec, PrimaryKeyCodec}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::{ColumnId, SequenceNumber, SequenceRange}; @@ -73,7 +72,7 @@ use crate::error::{ ComputeArrowSnafu, ComputeVectorSnafu, ConvertVectorSnafu, DecodeSnafu, InvalidBatchSnafu, Result, }; -use crate::memtable::{BoxedBatchIterator, BoxedRecordBatchIterator}; +use crate::memtable::BoxedRecordBatchIterator; pub(crate) fn timestamp_array_to_i64_slice(arr: &ArrayRef) -> &[i64] { use datatypes::arrow::array::{ @@ -119,6 +118,8 @@ pub struct Batch { primary_key: Vec, /// Possibly decoded `primary_key` values. Some places would decode it in advance. pk_values: Option, + /// Lazily decoded Dense columns, shared by consumers such as index creators. + dense_pk_cache: Option>, /// Timestamps of rows, should be sorted and not null. timestamps: VectorRef, /// Sequences of rows @@ -135,6 +136,13 @@ pub struct Batch { fields_idx: Option>, } +/// Only used when no fully decoded primary key was supplied by the reader. +#[derive(Debug, PartialEq, Clone)] +struct DensePkCache { + offsets: Vec, + values: Vec>, +} + impl Batch { /// Creates a new batch. pub fn new( @@ -173,12 +181,14 @@ impl Batch { /// Sets possibly decoded primary-key values. pub fn set_pk_values(&mut self, pk_values: CompositeValues) { self.pk_values = Some(pk_values); + self.dense_pk_cache = None; } /// Removes possibly decoded primary-key values. For testing only. #[cfg(any(test, feature = "test"))] pub fn remove_pk_values(&mut self) { self.pk_values = None; + self.dense_pk_cache = None; } /// Returns fields in the batch. @@ -214,6 +224,7 @@ impl Batch { Self { primary_key: vec![], pk_values: None, + dense_pk_cache: None, timestamps: Arc::new(TimestampMillisecondVectorBuilder::with_capacity(0).finish()), sequences: Arc::new(UInt64VectorBuilder::with_capacity(0).finish()), op_types: Arc::new(UInt8VectorBuilder::with_capacity(0).finish()), @@ -269,6 +280,7 @@ impl Batch { /// Be sure to update that field as well. pub fn set_primary_key(&mut self, primary_key: Vec) { self.primary_key = primary_key; + self.dense_pk_cache = None; } /// Slice the batch, returning a new batch. @@ -290,6 +302,7 @@ impl Batch { // this becomes a bottleneck. primary_key: self.primary_key.clone(), pk_values: self.pk_values.clone(), + dense_pk_cache: self.dense_pk_cache.clone(), timestamps: self.timestamps.slice(offset, length), sequences: Arc::new(self.sequences.get_slice(offset, length)), op_types: Arc::new(self.op_types.get_slice(offset, length)), @@ -692,89 +705,26 @@ impl Batch { self.sequences.get_data(index).unwrap() } - /// Checks the batch is monotonic by timestamps. - #[cfg(debug_assertions)] - #[allow(dead_code)] - pub(crate) fn check_monotonic(&self) -> Result<(), String> { - use std::cmp::Ordering; - if self.timestamps_native().is_none() { - return Ok(()); - } - - let timestamps = self.timestamps_native().unwrap(); - let sequences = self.sequences.as_arrow().values(); - for (i, window) in timestamps.windows(2).enumerate() { - let current = window[0]; - let next = window[1]; - let current_sequence = sequences[i]; - let next_sequence = sequences[i + 1]; - match current.cmp(&next) { - Ordering::Less => { - // The current timestamp is less than the next timestamp. - continue; - } - Ordering::Equal => { - // The current timestamp is equal to the next timestamp. - if current_sequence < next_sequence { - return Err(format!( - "sequence are not monotonic: ts {} == {} but current sequence {} < {}, index: {}", - current, next, current_sequence, next_sequence, i - )); - } - } - Ordering::Greater => { - // The current timestamp is greater than the next timestamp. - return Err(format!( - "timestamps are not monotonic: {} > {}, index: {}", - current, next, i - )); - } + /// Prepares a shared full-key cache when all Dense PK columns are needed. + /// Explicitly supplied decoded/defaulted values take precedence. + pub(crate) fn ensure_dense_pk_decoded(&mut self, codec: &DensePrimaryKeyCodec) -> Result<()> { + if self.pk_values.is_none() { + // A fallible iterator has no nonzero lower size hint. Reserve the + // known field count instead of growing its collected Vec per key. + let mut values = Vec::with_capacity(codec.num_fields()); + for value in codec.decode_dense_iter(&self.primary_key) { + values.push(value.context(DecodeSnafu)?); } + self.set_pk_values(CompositeValues::Dense(values)); } - Ok(()) } - /// Returns Ok if the given batch is behind the current batch. - #[cfg(debug_assertions)] - #[allow(dead_code)] - pub(crate) fn check_next_batch(&self, other: &Batch) -> Result<(), String> { - // Checks the primary key - if self.primary_key() < other.primary_key() { - return Ok(()); - } - if self.primary_key() > other.primary_key() { - return Err(format!( - "primary key is not monotonic: {:?} > {:?}", - self.primary_key(), - other.primary_key() - )); - } - // Checks the timestamp. - if self.last_timestamp() < other.first_timestamp() { - return Ok(()); - } - if self.last_timestamp() > other.first_timestamp() { - return Err(format!( - "timestamps are not monotonic: {:?} > {:?}", - self.last_timestamp(), - other.first_timestamp() - )); - } - // Checks the sequence. - if self.last_sequence() >= other.first_sequence() { - return Ok(()); - } - Err(format!( - "sequences are not monotonic: {:?} < {:?}", - self.last_sequence(), - other.first_sequence() - )) - } - /// Returns the value of the column in the primary key. /// - /// Lazily decodes the primary key and caches the result. + /// Reuses predecoded values when available. Otherwise Dense keys decode only + /// the requested column, sharing offsets and values across callers; Sparse + /// keys cache the full decode. pub fn pk_col_value( &mut self, codec: &dyn PrimaryKeyCodec, @@ -782,6 +732,24 @@ impl Batch { column_id: ColumnId, ) -> Result> { if self.pk_values.is_none() { + if let Some(codec) = codec.as_dense() { + if col_idx_in_pk >= codec.num_fields() { + return Ok(None); + } + let cache = self.dense_pk_cache.get_or_insert_with(|| { + Box::new(DensePkCache { + offsets: Vec::new(), + values: vec![None; codec.num_fields()], + }) + }); + if cache.values[col_idx_in_pk].is_none() { + let value = codec + .decode_value_at(&self.primary_key, col_idx_in_pk, &mut cache.offsets) + .context(DecodeSnafu)?; + cache.values[col_idx_in_pk] = Some(value); + } + return Ok(cache.values[col_idx_in_pk].as_ref()); + } self.pk_values = Some(codec.decode(&self.primary_key).context(DecodeSnafu)?); } @@ -814,112 +782,6 @@ impl Batch { } } -/// A struct to check the batch is monotonic. -#[cfg(debug_assertions)] -#[derive(Default)] -#[allow(dead_code)] -pub(crate) struct BatchChecker { - last_batch: Option, - start: Option, - end: Option, -} - -#[cfg(debug_assertions)] -#[allow(dead_code)] -impl BatchChecker { - /// Attaches the given start timestamp to the checker. - pub(crate) fn with_start(mut self, start: Option) -> Self { - self.start = start; - self - } - - /// Attaches the given end timestamp to the checker. - pub(crate) fn with_end(mut self, end: Option) -> Self { - self.end = end; - self - } - - /// Returns true if the given batch is monotonic and behind - /// the last batch. - pub(crate) fn check_monotonic(&mut self, batch: &Batch) -> Result<(), String> { - batch.check_monotonic()?; - - if let (Some(start), Some(first)) = (self.start, batch.first_timestamp()) - && start > first - { - return Err(format!( - "batch's first timestamp is before the start timestamp: {:?} > {:?}", - start, first - )); - } - if let (Some(end), Some(last)) = (self.end, batch.last_timestamp()) - && end <= last - { - return Err(format!( - "batch's last timestamp is after the end timestamp: {:?} <= {:?}", - end, last - )); - } - - // Checks the batch is behind the last batch. - // Then Updates the last batch. - let res = self - .last_batch - .as_ref() - .map(|last| last.check_next_batch(batch)) - .unwrap_or(Ok(())); - self.last_batch = Some(batch.clone()); - res - } - - /// Formats current batch and last batch for debug. - pub(crate) fn format_batch(&self, batch: &Batch) -> String { - use std::fmt::Write; - - let mut message = String::new(); - if let Some(last) = &self.last_batch { - write!( - message, - "last_pk: {:?}, last_ts: {:?}, last_seq: {:?}, ", - last.primary_key(), - last.last_timestamp(), - last.last_sequence() - ) - .unwrap(); - } - write!( - message, - "batch_pk: {:?}, batch_ts: {:?}, batch_seq: {:?}", - batch.primary_key(), - batch.timestamps(), - batch.sequences() - ) - .unwrap(); - - message - } - - /// Checks batches from the part range are monotonic. Otherwise, panics. - pub(crate) fn ensure_part_range_batch( - &mut self, - scanner: &str, - region_id: store_api::storage::RegionId, - partition: usize, - part_range: store_api::region_engine::PartitionRange, - batch: &Batch, - ) { - if let Err(e) = self.check_monotonic(batch) { - let err_msg = format!( - "{}: batch is not sorted, {}, region_id: {}, partition: {}, part_range: {:?}", - scanner, e, region_id, partition, part_range, - ); - common_telemetry::error!("{err_msg}, {}", self.format_batch(batch)); - // Only print the number of row in the panic message. - panic!("{err_msg}, batch rows: {}", batch.num_rows()); - } - } -} - /// Len of timestamp in arrow row format. const TIMESTAMP_KEY_LEN: usize = 9; @@ -1098,6 +960,7 @@ impl BatchBuilder { Ok(Batch { primary_key: self.primary_key, pk_values: None, + dense_pk_cache: None, timestamps, sequences, op_types, @@ -1119,29 +982,6 @@ impl From for BatchBuilder { } } -/// Async [Batch] reader and iterator wrapper. -/// -/// This is the data source for SST writers or internal readers. -pub enum Source { - /// Source from a [BoxedBatchReader]. - Reader(BoxedBatchReader), - /// Source from a [BoxedBatchIterator]. - Iter(BoxedBatchIterator), - /// Source from a [BoxedBatchStream]. - Stream(BoxedBatchStream), -} - -impl Source { - /// Returns next [Batch] from this data source. - pub async fn next_batch(&mut self) -> Result> { - match self { - Source::Reader(reader) => reader.next_batch().await, - Source::Iter(iter) => iter.next().transpose(), - Source::Stream(stream) => stream.try_next().await, - } - } -} - /// Async [RecordBatch] reader and iterator wrapper for flat format. pub struct FlatSource { schema: SchemaRef, @@ -1199,37 +1039,9 @@ impl FlatSourceInner { } } -/// Async batch reader. -/// -/// The reader must guarantee [Batch]es returned by it have the same schema. -#[async_trait] -pub trait BatchReader: Send { - /// Fetch next [Batch]. - /// - /// Returns `Ok(None)` when the reader has reached its end and calling `next_batch()` - /// again won't return batch again. - /// - /// If `Err` is returned, caller should not call this method again, the implementor - /// may or may not panic in such case. - async fn next_batch(&mut self) -> Result>; -} - -/// Pointer to [BatchReader]. -pub type BoxedBatchReader = Box; - -/// Pointer to a stream that yields [Batch]. -pub type BoxedBatchStream = BoxStream<'static, Result>; - /// Pointer to a stream that yields [RecordBatch]. pub type BoxedRecordBatchStream = BoxStream<'static, Result>; -#[async_trait::async_trait] -impl BatchReader for Box { - async fn next_batch(&mut self) -> Result> { - (**self).next_batch().await - } -} - /// Local metrics for scanners. #[derive(Debug, Default)] pub(crate) struct ScannerMetrics { @@ -1254,6 +1066,79 @@ mod tests { use crate::error::Error; use crate::test_util::new_batch_builder; + #[test] + fn dense_pk_columns_are_lazy_and_reset_with_the_key() { + use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; + + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (7, SortField::new(ConcreteDataType::string_datatype())), + (3, SortField::new(ConcreteDataType::int64_datatype())), + (9, SortField::new(ConcreteDataType::string_datatype())), + ]); + let values = [ + Value::from("中文\0abcdefgh"), + Value::Int64(-42), + Value::Null, + ]; + let key = codec + .encode(values.iter().map(Value::as_value_ref)) + .unwrap(); + let mut batch = new_batch_without_fields(&[1, 2], &[1, 1], &[OpType::Put, OpType::Put]); + batch.set_primary_key(key); + for pos in [2, 0, 1, 2] { + assert_eq!( + batch.pk_col_value(&codec, pos, [7, 3, 9][pos]).unwrap(), + Some(&values[pos]) + ); + } + assert!(batch.pk_col_value(&codec, 3, 99).unwrap().is_none()); + let mut sliced = batch.slice(1, 1); + assert_eq!(sliced.pk_col_value(&codec, 0, 7).unwrap(), Some(&values[0])); + + // A corrupt unrequested suffix must not force whole-key decoding. + batch.set_primary_key(vec![0, 1]); + assert_eq!( + batch.pk_col_value(&codec, 0, 7).unwrap(), + Some(&Value::Null) + ); + assert!(batch.pk_col_value(&codec, 1, 3).is_err()); + batch.set_primary_key( + codec + .encode( + [ + ValueRef::String(""), + ValueRef::Int64(8), + ValueRef::String("new"), + ] + .into_iter(), + ) + .unwrap(), + ); + assert_eq!( + batch.pk_col_value(&codec, 2, 9).unwrap(), + Some(&Value::from("new")) + ); + assert_eq!( + batch.pk_col_value(&codec, 0, 7).unwrap(), + Some(&Value::from("")) + ); + + // Schema compatibility may supply already decoded/defaulted values. + batch.set_pk_values(CompositeValues::Dense(vec![(7, Value::from("default"))])); + batch.ensure_dense_pk_decoded(&codec).unwrap(); + assert_eq!( + batch.pk_col_value(&codec, 0, 7).unwrap(), + Some(&Value::from("default")) + ); + assert!(batch.pk_col_value(&codec, 1, 3).unwrap().is_none()); + batch.remove_pk_values(); + batch.ensure_dense_pk_decoded(&codec).unwrap(); + assert_eq!( + batch.pk_col_value(&codec, 1, 3).unwrap(), + Some(&Value::Int64(8)) + ); + } + fn new_batch( timestamps: &[i64], sequences: &[u64], diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 0e566b7e078..acc05f0e9d2 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use api::v1::SemanticType; use datatypes::arrow::array::{ - Array, ArrayRef, BinaryArray, BinaryBuilder, DictionaryArray, UInt32Array, + Array, ArrayRef, BinaryArray, BinaryBuilder, DictionaryArray, UInt32Array, new_null_array, }; use datatypes::arrow::compute::{TakeOptions, take}; use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef}; @@ -27,7 +27,7 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::prelude::DataType; use datatypes::value::Value; -use datatypes::vectors::VectorRef; +use datatypes::vectors::{Helper, VectorRef}; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec, build_primary_key_codec_with_fields, @@ -192,20 +192,32 @@ impl FlatCompatBatch { }); } else { // Create a default vector with 1 element for that column. - let default_vector = expect_column - .column_schema - .create_default_vector(1) - .context(CreateDefaultSnafu { - region_id: expect_metadata.region_id, - column: &expect_column.column_schema.name, - })? - .with_context(|| CompatReaderSnafu { - region_id: expect_metadata.region_id, - reason: format!( - "column {} does not have a default value to read", - expect_column.column_schema.name - ), - })?; + let default_vector = if expect_data_type.is_json2() { + // JSON2 columns do not support defaults, so missing columns use NULL. + // `expect_data_type` may be a query-concretized JSON2 type, + // while region metadata still carries its unconcretized type. + // The default vector must match the target Arrow field above. + Helper::try_into_vector(new_null_array(&expect_data_type.as_arrow_type(), 1)) + .context(CreateDefaultSnafu { + region_id: expect_metadata.region_id, + column: &expect_column.column_schema.name, + })? + } else { + expect_column + .column_schema + .create_default_vector(1) + .context(CreateDefaultSnafu { + region_id: expect_metadata.region_id, + column: &expect_column.column_schema.name, + })? + .with_context(|| CompatReaderSnafu { + region_id: expect_metadata.region_id, + reason: format!( + "column {} does not have a default value to read", + expect_column.column_schema.name + ), + })? + }; index_or_defaults.push(IndexOrDefault::DefaultValue { default_vector, semantic_type: expect_column.semantic_type, diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index c526bfcaa86..2bafe02dbbe 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -16,7 +16,6 @@ use std::sync::Arc; -use async_trait::async_trait; use datatypes::arrow::array::{Array, BinaryArray}; use datatypes::arrow::compute::concat_batches; use datatypes::arrow::record_batch::RecordBatch; @@ -30,98 +29,13 @@ use crate::cache::{ }; use crate::error::{ComputeArrowSnafu, Result}; use crate::read::read_columns::JsonTargetTypes; -use crate::read::{ - Batch, BatchReader, BoxedBatchReader, BoxedRecordBatchStream, timestamp_array_to_i64_slice, -}; +use crate::read::{BoxedRecordBatchStream, timestamp_array_to_i64_slice}; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::parquet::format::{PrimaryKeyArray, primary_key_offsets}; use crate::sst::parquet::read_columns::ParquetReadColumns; use crate::sst::parquet::reader::FlatRowGroupReader; -/// Reader to keep the last row for each time series. -/// It assumes that batches from the input reader are -/// - sorted -/// - all deleted rows has been filtered. -/// - not empty -/// -/// This reader is different from the [MergeMode](crate::region::options::MergeMode) as -/// it focus on time series (the same key). -#[allow(dead_code)] -pub(crate) struct LastRowReader { - /// Inner reader. - reader: BoxedBatchReader, - /// The last batch pending to return. - selector: LastRowSelector, -} - -#[allow(dead_code)] -impl LastRowReader { - /// Creates a new `LastRowReader`. - pub(crate) fn new(reader: BoxedBatchReader) -> Self { - Self { - reader, - selector: LastRowSelector::default(), - } - } - - /// Returns the last row of the next key. - pub(crate) async fn next_last_row(&mut self) -> Result> { - while let Some(batch) = self.reader.next_batch().await? { - if let Some(yielded) = self.selector.on_next(batch) { - return Ok(Some(yielded)); - } - } - Ok(self.selector.finish()) - } -} - -#[async_trait] -impl BatchReader for LastRowReader { - async fn next_batch(&mut self) -> Result> { - self.next_last_row().await - } -} - -/// Common struct that selects only the last row of each time series. -#[derive(Default)] -pub struct LastRowSelector { - last_batch: Option, -} - -impl LastRowSelector { - /// Handles next batch. Return the yielding batch if present. - pub fn on_next(&mut self, batch: Batch) -> Option { - if let Some(last) = &self.last_batch { - if last.primary_key() == batch.primary_key() { - // Same key, update last batch. - self.last_batch = Some(batch); - None - } else { - // Different key, return the last row in `last` and update `last_batch` by - // current batch. - debug_assert!(!last.is_empty()); - let last_row = last.slice(last.num_rows() - 1, 1); - self.last_batch = Some(batch); - Some(last_row) - } - } else { - self.last_batch = Some(batch); - None - } - } - - /// Finishes the selector and returns the pending batch if any. - pub fn finish(&mut self) -> Option { - if let Some(last) = self.last_batch.take() { - // This is the last key. - let last_row = last.slice(last.num_rows() - 1, 1); - return Some(last_row); - } - None - } -} - /// Cached last row reader for flat format row group. /// If the last rows are already cached (as flat `RecordBatch`), returns cached values. /// Otherwise, reads from the row group, selects last rows, and updates the cache. @@ -514,7 +428,6 @@ fn last_timestamp_start(ts_values: &[i64], range_start: usize, range_end: usize) mod tests { use std::sync::Arc; - use api::v1::OpType; use datatypes::arrow::array::{ ArrayRef, BinaryDictionaryBuilder, Int64Array, TimestampMillisecondArray, UInt8Array, UInt64Array, @@ -523,72 +436,6 @@ mod tests { use datatypes::arrow::record_batch::RecordBatch; use super::*; - use crate::test_util::{VecBatchReader, check_reader_result, new_batch}; - - #[tokio::test] - async fn test_last_row_one_batch() { - let input = [new_batch( - b"k1", - &[1, 2], - &[11, 11], - &[OpType::Put, OpType::Put], - &[21, 22], - )]; - let reader = VecBatchReader::new(&input); - let mut reader = LastRowReader::new(Box::new(reader)); - check_reader_result( - &mut reader, - &[new_batch(b"k1", &[2], &[11], &[OpType::Put], &[22])], - ) - .await; - - // Only one row. - let input = [new_batch(b"k1", &[1], &[11], &[OpType::Put], &[21])]; - let reader = VecBatchReader::new(&input); - let mut reader = LastRowReader::new(Box::new(reader)); - check_reader_result( - &mut reader, - &[new_batch(b"k1", &[1], &[11], &[OpType::Put], &[21])], - ) - .await; - } - - #[tokio::test] - async fn test_last_row_multi_batch() { - let input = [ - new_batch( - b"k1", - &[1, 2], - &[11, 11], - &[OpType::Put, OpType::Put], - &[21, 22], - ), - new_batch( - b"k1", - &[3, 4], - &[11, 11], - &[OpType::Put, OpType::Put], - &[23, 24], - ), - new_batch( - b"k2", - &[1, 2], - &[11, 11], - &[OpType::Put, OpType::Put], - &[31, 32], - ), - ]; - let reader = VecBatchReader::new(&input); - let mut reader = LastRowReader::new(Box::new(reader)); - check_reader_result( - &mut reader, - &[ - new_batch(b"k1", &[4], &[11], &[OpType::Put], &[24]), - new_batch(b"k2", &[2], &[11], &[OpType::Put], &[32]), - ], - ) - .await; - } /// Helper to build a flat format RecordBatch for testing. fn new_flat_batch(primary_keys: &[&[u8]], timestamps: &[i64], fields: &[i64]) -> RecordBatch { diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs index 19c487ed0b8..ef09f5f48ab 100644 --- a/src/mito2/src/read/pruner.rs +++ b/src/mito2/src/read/pruner.rs @@ -350,12 +350,14 @@ impl Pruner { enable_predicate_prefilter, }); - // Spawn worker tasks with their receivers + // Keep pruning and prefetching on the runtime of the originating workload. for (worker_id, rx) in receivers.into_iter().enumerate() { - let inner_clone = inner.clone(); - common_runtime::spawn_query(async move { - Self::worker_loop(worker_id, rx, inner_clone).await; - }); + let worker = Self::worker_loop(worker_id, rx, inner.clone()); + if inner.stream_ctx.input.compaction { + common_runtime::spawn_compact(worker); + } else { + common_runtime::spawn_query(worker); + } } Self { diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 1d0991a6bfb..0f9cea941b3 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -29,10 +29,12 @@ use common_telemetry::tracing::Instrument; use common_telemetry::{debug, error, tracing, warn}; use common_time::range::TimestampRange; use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::physical_plan::expressions::DynamicFilterPhysicalExpr; use datafusion_common::pruning::PruningStatistics; -use datafusion_common::{Column, ScalarValue}; +use datafusion_common::{Column, ScalarValue, ToDFSchema}; use datafusion_expr::Expr; +use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use datatypes::extension::json::is_json2_extension_type; @@ -57,8 +59,8 @@ use crate::access_layer::AccessLayerRef; use crate::cache::CacheStrategy; use crate::config::DEFAULT_MAX_CONCURRENT_SCAN_FILES; use crate::error::{ - InvalidPartitionExprSnafu, InvalidRequestSnafu, RegionSequenceDomainBrokenSnafu, Result, - SequenceRangeUnsupportedSnafu, + EvalPartitionFilterSnafu, InvalidPartitionExprSnafu, InvalidRequestSnafu, + RegionSequenceDomainBrokenSnafu, Result, SequenceRangeUnsupportedSnafu, }; #[cfg(feature = "enterprise")] use crate::extension::{BoxedExtensionRange, BoxedExtensionRangeProvider}; @@ -85,16 +87,11 @@ use crate::sst::index::fulltext_index::applier::FulltextIndexApplierRef; use crate::sst::index::fulltext_index::applier::builder::FulltextIndexApplierBuilder; use crate::sst::index::inverted_index::applier::InvertedIndexApplierRef; use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; -#[cfg(feature = "vector_index")] -use crate::sst::index::vector_index::applier::{VectorIndexApplier, VectorIndexApplierRef}; use crate::sst::parquet::Json2RewriteTargets; use crate::sst::parquet::file_range::PreFilterMode; use crate::sst::parquet::reader::ReaderMetrics; use crate::sst::primary_key::PrimaryKeyRangeMapper; -#[cfg(feature = "vector_index")] -const VECTOR_INDEX_OVERFETCH_MULTIPLIER: usize = 2; - /// A scanner scans a region and returns a [SendableRecordBatchStream]. pub(crate) enum Scanner { /// Sequential scan. @@ -572,16 +569,6 @@ impl ScanRegion { self.build_fulltext_index_applier(&non_field_filters), self.build_fulltext_index_applier(&field_filters), ]; - #[cfg(feature = "vector_index")] - let vector_index_applier = self.build_vector_index_applier(); - #[cfg(feature = "vector_index")] - let vector_index_k = self.request.vector_search.as_ref().map(|search| { - if self.request.filters.is_empty() { - search.k - } else { - search.k.saturating_mul(VECTOR_INDEX_OVERFETCH_MULTIPLIER) - } - }); let input = ScanInput::builder(self.access_layer, mapper) .with_series_index(self.series_index) @@ -614,11 +601,6 @@ impl ScanRegion { ) .with_sequence_range(sequence_range) .with_query_stat_counters(self.query_stat_counters); - #[cfg(feature = "vector_index")] - let input = input - .with_vector_index_applier(vector_index_applier) - .with_vector_index_k(vector_index_k); - #[cfg(feature = "enterprise")] let input = if !self.request.skip_sst_files && let Some(provider) = self.extension_range_provider @@ -906,31 +888,6 @@ impl ScanRegion { .flatten() .map(Arc::new) } - - /// Build the vector index applier from vector search request. - #[cfg(feature = "vector_index")] - fn build_vector_index_applier(&self) -> Option { - let vector_search = self.request.vector_search.as_ref()?; - - let file_cache = self.cache_strategy.write_cache().map(|w| w.file_cache()); - let puffin_metadata_cache = self.cache_strategy.puffin_metadata_cache().cloned(); - let vector_index_cache = self.cache_strategy.vector_index_cache().cloned(); - - let applier = VectorIndexApplier::new( - self.access_layer.table_dir().to_string(), - self.access_layer.path_type(), - self.access_layer.object_store().clone(), - self.access_layer.puffin_manager_factory().clone(), - vector_search.column_id, - vector_search.query_vector.clone(), - vector_search.metric, - ) - .with_file_cache(file_cache) - .with_puffin_metadata_cache(puffin_metadata_cache) - .with_vector_index_cache(vector_index_cache); - - Some(Arc::new(applier)) - } } /// Returns true if the time range of a SST `file` matches the `predicate`. @@ -995,12 +952,6 @@ pub struct ScanInput { inverted_index_appliers: [Option; 2], bloom_filter_index_appliers: [Option; 2], fulltext_index_appliers: [Option; 2], - /// Vector index applier for KNN search. - #[cfg(feature = "vector_index")] - pub(crate) vector_index_applier: Option, - /// Over-fetched k for vector index scan. - #[cfg(feature = "vector_index")] - pub(crate) vector_index_k: Option, /// Start time of the query. pub(crate) query_start: Option, /// The region is using append mode. @@ -1079,10 +1030,6 @@ impl ScanInput { inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], fulltext_index_appliers: [None, None], - #[cfg(feature = "vector_index")] - vector_index_applier: None, - #[cfg(feature = "vector_index")] - vector_index_k: None, query_start: None, append_mode: false, filter_deleted: true, @@ -1365,25 +1312,6 @@ impl ScanInputBuilder { self } - /// Sets vector index applier for KNN search. - #[cfg(feature = "vector_index")] - #[must_use] - pub(crate) fn with_vector_index_applier( - mut self, - applier: Option, - ) -> Self { - self.input.vector_index_applier = applier; - self - } - - /// Sets over-fetched k for vector index scan. - #[cfg(feature = "vector_index")] - #[must_use] - pub(crate) fn with_vector_index_k(mut self, k: Option) -> Self { - self.input.vector_index_k = k; - self - } - /// Sets start time of the query. #[must_use] pub(crate) fn with_start_time(mut self, now: Option) -> Self { @@ -1632,13 +1560,6 @@ impl ScanInput { } else { reader }; - #[cfg(feature = "vector_index")] - let reader = { - let mut reader = reader; - reader = - reader.vector_index_applier(self.vector_index_applier.clone(), self.vector_index_k); - reader - }; let res = reader .expected_metadata(Some(self.mapper.metadata().clone())) .compaction(self.compaction) @@ -1744,7 +1665,15 @@ impl ScanInput { }; match maybe_batch { Some(Ok(batch)) => { - let _ = sender.send(Ok(batch)).await; + // The receiver is gone when the query is cancelled or finishes early, + // so stop reading the source. + if let Err(e) = sender.send(Ok(batch)).await { + debug!( + "Stop parallel scan task, receiver dropped, region_id: {}, error: {}", + region_id, e + ); + break; + } } Some(Err(e)) => { let _ = sender.send(Err(e)).await; @@ -2197,10 +2126,6 @@ impl StreamContext { write!(f, ", \"dyn_filters\": {:?}", dyn_filters)?; } } - #[cfg(feature = "vector_index")] - if let Some(vector_index_k) = self.input.vector_index_k { - write!(f, ", \"vector_index_k\": {}", vector_index_k)?; - } if !self.input.files.is_empty() { write!(f, ", \"files\": ")?; f.debug_list() @@ -2274,6 +2199,22 @@ impl PredicateGroup { expr: expr_json.clone(), })?; + // Persisted partition bounds retain their original types after schema changes. + let schema = metadata + .schema + .arrow_schema() + .clone() + .to_dfschema_ref() + .context(EvalPartitionFilterSnafu)?; + let simplifier = ExprSimplifier::new( + SimplifyContext::builder() + .with_schema(schema.clone()) + .build(), + ); + let logical_expr = simplifier + .coerce(logical_expr, &schema) + .and_then(|expr| simplifier.simplify(expr)) + .context(EvalPartitionFilterSnafu)?; combined_exprs.push(logical_expr); region_partition_expr = Some(expr); } @@ -2416,6 +2357,75 @@ mod tests { .with_files(vec![file]) } + #[test] + fn test_partition_time_filter_after_unit_widening() { + use datatypes::arrow::array::TimestampMicrosecondArray; + + let bound = Value::Timestamp(Timestamp::new_millisecond(1)); + for (expr, expected) in [ + ( + partition_col("ts").lt(bound.clone()), + vec![true, false, false], + ), + (partition_col("ts").gt_eq(bound), vec![false, true, true]), + ] { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_microsecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }); + builder.primary_key(vec![]); + builder.partition_expr_json(Some(expr.as_json_str().unwrap())); + let metadata = builder.build().unwrap(); + let group = PredicateGroup::new(&metadata, &[]).unwrap(); + let filters = group.time_filters().unwrap(); + let input: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![999, 1000, 1001])); + assert_eq!(filters.len(), 1); + assert_eq!( + filters[0].evaluate_array(&input).unwrap(), + datafusion_common::arrow::buffer::BooleanBuffer::from(expected), + ); + } + } + + #[test] + fn test_partition_time_filter_overflow_after_unit_widening() { + // Year 3000 fits milliseconds but overflows nanoseconds. + let bound = Value::Timestamp(Timestamp::new_millisecond(32_503_680_000_000)); + let expr = partition_col("ts").lt(bound); + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }); + builder.primary_key(vec![]); + builder.partition_expr_json(Some(expr.as_json_str().unwrap())); + let metadata = builder.build().unwrap(); + + let error = PredicateGroup::new(&metadata, &[]) + .err() + .expect("an overflowing partition bound must fail scan construction"); + let crate::error::Error::EvalPartitionFilter { error, .. } = error else { + panic!("unexpected error: {error:?}"); + }; + assert!( + error + .to_string() + .contains("converted value exceeds the representable i64 range"), + "{error}" + ); + } + #[tokio::test] async fn test_total_rows_is_exact_after_partition_filter() { let expr = partition_col("k0").gt_eq(Value::String("foo".into())); @@ -2878,12 +2888,12 @@ mod tests { }) .primary_key(vec![0]); let metadata = Arc::new(builder.build().unwrap()); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let version = Arc::new( crate::region::version::VersionBuilder::new(metadata.clone(), mutable).build(), ); @@ -3157,12 +3167,12 @@ mod tests { ..Default::default() }; let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let version = Arc::new( crate::region::version::VersionBuilder::new(metadata, mutable) .options(crate::region::options::RegionOptions { @@ -3197,12 +3207,12 @@ mod tests { ..Default::default() }; let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let target_region_id = metadata.region_id; // Files at or below the cursor are excluded. Among the remaining files, @@ -3214,23 +3224,25 @@ mod tests { (None, false, true, false), (None, true, true, true), ] { - let version = - crate::region::version::VersionBuilder::new(metadata.clone(), mutable.clone()) - .options(crate::region::options::RegionOptions { - preserve_row_sequence: true, - ..Default::default() - }) - .add_files( - new_noop_file_purger(), - [FileMeta { - region_id: target_region_id, - sequence, - preserve_row_sequence: marked, - ..Default::default() - }] - .into_iter(), - ) - .build(); + let version = crate::region::version::VersionBuilder::new( + metadata.clone(), + mutable.new_with_part_duration(None, None), + ) + .options(crate::region::options::RegionOptions { + preserve_row_sequence: true, + ..Default::default() + }) + .add_files( + new_noop_file_purger(), + [FileMeta { + region_id: target_region_id, + sequence, + preserve_row_sequence: marked, + ..Default::default() + }] + .into_iter(), + ) + .build(); let (files, range) = exact_sequence_range(&request, &version).unwrap(); assert_eq!(expected_files, !files.is_empty()); assert_eq!(expected_range, range.is_some()); @@ -3245,24 +3257,26 @@ mod tests { (None, true, 1, true), ] { let file_id = store_api::storage::FileId::random(); - let version = - crate::region::version::VersionBuilder::new(metadata.clone(), mutable.clone()) - .options(crate::region::options::RegionOptions { - preserve_row_sequence: true, - ..Default::default() - }) - .add_files( - new_noop_file_purger(), - [FileMeta { - region_id: RegionId::new(1, 2), - file_id, - sequence, - preserve_row_sequence: marked, - ..Default::default() - }] - .into_iter(), - ) - .build(); + let version = crate::region::version::VersionBuilder::new( + metadata.clone(), + mutable.new_with_part_duration(None, None), + ) + .options(crate::region::options::RegionOptions { + preserve_row_sequence: true, + ..Default::default() + }) + .add_files( + new_noop_file_purger(), + [FileMeta { + region_id: RegionId::new(1, 2), + file_id, + sequence, + preserve_row_sequence: marked, + ..Default::default() + }] + .into_iter(), + ) + .build(); let result = exact_sequence_range(&request, &version); if expected_error { assert!(result.is_err()); @@ -3295,4 +3309,59 @@ mod tests { .build(); assert!(exact_sequence_range(&request, &version).is_err()); } + + #[tokio::test] + async fn test_parallel_flat_source_stops_after_receiver_dropped() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use datatypes::arrow::array::Int64Array; + + const SOURCE_BATCHES: usize = 10_000; + + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap(); + let env = SchedulerEnv::new().await; + let input = ScanInput::builder(env.access_layer.clone(), mapper).build(); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "v", + ArrowDataType::Int64, + false, + )])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap(); + let pulled = Arc::new(AtomicUsize::new(0)); + let mut released = Vec::new(); + let sources = (0..2) + .map(|_| { + let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + released.push(release_rx); + let batch = batch.clone(); + let pulled = pulled.clone(); + futures::stream::iter(0..SOURCE_BATCHES) + .map(move |_| { + // Dropped together with the source stream. + let _ = &release_tx; + pulled.fetch_add(1, Ordering::Relaxed); + Ok(batch.clone()) + }) + .boxed() + }) + .collect(); + + let streams = input + .create_parallel_flat_sources(sources, Arc::new(Semaphore::new(2)), 1) + .unwrap(); + drop(streams); + + for release_rx in released { + tokio::time::timeout(std::time::Duration::from_secs(10), release_rx) + .await + .unwrap() + .unwrap_err(); + } + // Each task may have read a few batches ahead into its channel before the + // receiver went away, but must not drain the rest of its source. + assert!(pulled.load(Ordering::Relaxed) < 100); + } } diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 2d93cbbbe03..c5eed8853ae 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -158,8 +158,6 @@ pub(crate) struct ScanMetricsSet { rg_minmax_filtered: usize, /// Number of row groups filtered by bloom filter index. rg_bloom_filtered: usize, - /// Number of row groups filtered by vector index. - rg_vector_filtered: usize, /// Number of rows in row group before filtering. rows_before_filter: usize, /// Number of rows in row group filtered by fulltext index. @@ -168,10 +166,6 @@ pub(crate) struct ScanMetricsSet { rows_inverted_filtered: usize, /// Number of rows in row group filtered by bloom filter index. rows_bloom_filtered: usize, - /// Number of rows filtered by vector index. - rows_vector_filtered: usize, - /// Number of rows selected by vector index. - rows_vector_selected: usize, /// Number of rows filtered by precise filter. rows_precise_filtered: usize, /// Number of index result cache hits for fulltext index. @@ -310,13 +304,10 @@ impl fmt::Debug for ScanMetricsSet { rg_inverted_filtered, rg_minmax_filtered, rg_bloom_filtered, - rg_vector_filtered, rows_before_filter, rows_fulltext_filtered, rows_inverted_filtered, rows_bloom_filtered, - rows_vector_filtered, - rows_vector_selected, rows_precise_filtered, fulltext_index_cache_hit, fulltext_index_cache_miss, @@ -409,9 +400,6 @@ impl fmt::Debug for ScanMetricsSet { if *rg_bloom_filtered > 0 { write!(f, ", \"rg_bloom_filtered\":{rg_bloom_filtered}")?; } - if *rg_vector_filtered > 0 { - write!(f, ", \"rg_vector_filtered\":{rg_vector_filtered}")?; - } if *rows_fulltext_filtered > 0 { write!(f, ", \"rows_fulltext_filtered\":{rows_fulltext_filtered}")?; } @@ -421,12 +409,6 @@ impl fmt::Debug for ScanMetricsSet { if *rows_bloom_filtered > 0 { write!(f, ", \"rows_bloom_filtered\":{rows_bloom_filtered}")?; } - if *rows_vector_filtered > 0 { - write!(f, ", \"rows_vector_filtered\":{rows_vector_filtered}")?; - } - if *rows_vector_selected > 0 { - write!(f, ", \"rows_vector_selected\":{rows_vector_selected}")?; - } if *rows_precise_filtered > 0 { write!(f, ", \"rows_precise_filtered\":{rows_precise_filtered}")?; } @@ -677,13 +659,10 @@ impl ScanMetricsSet { rg_inverted_filtered, rg_minmax_filtered, rg_bloom_filtered, - rg_vector_filtered, rows_total, rows_fulltext_filtered, rows_inverted_filtered, rows_bloom_filtered, - rows_vector_filtered, - rows_vector_selected, rows_precise_filtered, fulltext_index_cache_hit, fulltext_index_cache_miss, @@ -721,14 +700,11 @@ impl ScanMetricsSet { self.rg_inverted_filtered += *rg_inverted_filtered; self.rg_minmax_filtered += *rg_minmax_filtered; self.rg_bloom_filtered += *rg_bloom_filtered; - self.rg_vector_filtered += *rg_vector_filtered; self.rows_before_filter += *rows_total; self.rows_fulltext_filtered += *rows_fulltext_filtered; self.rows_inverted_filtered += *rows_inverted_filtered; self.rows_bloom_filtered += *rows_bloom_filtered; - self.rows_vector_filtered += *rows_vector_filtered; - self.rows_vector_selected += *rows_vector_selected; self.rows_precise_filtered += *rows_precise_filtered; self.fulltext_index_cache_hit += *fulltext_index_cache_hit; @@ -861,10 +837,6 @@ impl ScanMetricsSet { READ_ROW_GROUPS_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rg_bloom_filtered as u64); - #[cfg(feature = "vector_index")] - READ_ROW_GROUPS_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rg_vector_filtered as u64); PRECISE_FILTER_ROWS_TOTAL .with_label_values(&["parquet"]) @@ -881,10 +853,6 @@ impl ScanMetricsSet { READ_ROWS_IN_ROW_GROUP_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rows_bloom_filtered as u64); - #[cfg(feature = "vector_index")] - READ_ROWS_IN_ROW_GROUP_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rows_vector_filtered as u64); } } @@ -1144,21 +1112,18 @@ impl PartitionMetrics { } /// Increments the total bytes added to the range cache. - #[allow(dead_code)] pub(crate) fn inc_range_cache_size(&self, size: usize) { let mut metrics = self.0.metrics.lock().unwrap(); metrics.range_cache_size += size; } /// Increments the range cache hit counter. - #[allow(dead_code)] pub(crate) fn inc_range_cache_hit(&self) { let mut metrics = self.0.metrics.lock().unwrap(); metrics.range_cache_hit += 1; } /// Increments the range cache miss counter. - #[allow(dead_code)] pub(crate) fn inc_range_cache_miss(&self) { let mut metrics = self.0.metrics.lock().unwrap(); metrics.range_cache_miss += 1; diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index f05f54e6662..3ba5512a7bd 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -778,6 +778,7 @@ mod tests { )); let store = env.access_layer.object_store().clone(); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_millisecond(0), bucket_end: Timestamp::new_millisecond(20), diff --git a/src/mito2/src/read/stream.rs b/src/mito2/src/read/stream.rs index 27ae0e8c3a1..0c882d374d2 100644 --- a/src/mito2/src/read/stream.rs +++ b/src/mito2/src/read/stream.rs @@ -43,7 +43,6 @@ pub type ScanBatchStream = BoxStream<'static, Result>; pub(crate) struct ConvertBatchStream { inner: ScanBatchStream, projection_mapper: Arc, - #[allow(dead_code)] cache_strategy: CacheStrategy, partition_metrics: PartitionMetrics, pending: VecDeque, diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index fb5d260fe6f..71c1192cbf2 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -58,7 +58,7 @@ use crate::error::{ use crate::manifest::action::RegionManifest; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::memtable::bulk::part::BulkPart; -use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef}; +use crate::memtable::time_partition::TimePartitions; use crate::memtable::{MemtableBuilderProvider, ensure_json2_not_use_time_series_memtable}; use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES}; use crate::read::series_candidate::is_sparse_metric_metadata; @@ -68,7 +68,7 @@ use crate::region::{ ManifestContext, ManifestStats, MitoRegion, MitoRegionRef, RegionLeaderState, RegionRoleState, RegionStats, }; -use crate::region_write_ctx::RegionWriteCtx; +use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::OptionOutputTx; use crate::schedule::scheduler::SchedulerRef; use crate::series_index::{IndexFilePurger, load_version_control}; @@ -410,12 +410,8 @@ impl RegionOpener { let memtable_builder = self.memtable_builder_provider.builder_for_options(&options); let part_duration = options.compaction.time_window(); // Initial memtable id is 0. - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - memtable_builder.clone(), - 0, - part_duration, - )); + let mutable = + TimePartitions::new(metadata.clone(), memtable_builder.clone(), 0, part_duration); debug!( "Create region {} with options: {:?}, default_flat_format: {}", @@ -587,12 +583,8 @@ impl RegionOpener { .time_window() .or(manifest.compaction_time_window); // Initial memtable id is 0. - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - memtable_builder.clone(), - 0, - part_duration, - )); + let mutable = + TimePartitions::new(metadata.clone(), memtable_builder.clone(), 0, part_duration); // Updates region options by manifest before creating version. let version_builder = version_builder_from_manifest( @@ -785,7 +777,7 @@ pub(crate) fn version_builder_from_manifest( manifest: &RegionManifest, metadata: RegionMetadataRef, file_purger: FilePurgerRef, - mutable: TimePartitionsRef, + mutable: TimePartitions, region_options: RegionOptions, ) -> VersionBuilder { VersionBuilder::new(metadata, mutable) @@ -961,6 +953,7 @@ where provider.clone(), // For WAL replay, we don't need to track the write bytes rate. None, + WriteSource::WalReplay, ); for mutation in entry.mutations { rows_replayed += mutation @@ -1009,11 +1002,9 @@ where // set next_entry_id and write to memtable. region_write_ctx.set_next_entry_id(last_entry_id + 1); - region_write_ctx.write_memtable().await; - region_write_ctx.write_bulk().await; // Publish the replayed sequences only after all rows (including bulk // parts) are installed, matching the write path ordering. - region_write_ctx.publish_sequence_and_entry_id(); + region_write_ctx.write_memtables().await; } // TODO(weny): We need to update `flushed_entry_id` in the region manifest @@ -1247,7 +1238,12 @@ async fn preload_parquet_meta_cache_for_files( let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); if let Some(metadata) = write_cache .file_cache() - .get_sst_meta_data(key, &mut cache_metrics, PageIndexPolicy::Optional) + .get_sst_meta_data( + key, + &mut cache_metrics, + PageIndexPolicy::Optional, + &common_runtime::global_runtime(), + ) .await { let decoded = metadata.decoded(); @@ -1297,6 +1293,7 @@ async fn preload_parquet_meta_cache_for_files( // instead of substituting the region's current schema. None, PageIndexPolicy::Optional, + &common_runtime::global_runtime(), ) .await { diff --git a/src/mito2/src/region/utils.rs b/src/mito2/src/region/utils.rs index 6917315757a..7dfaad72aa7 100644 --- a/src/mito2/src/region/utils.rs +++ b/src/mito2/src/region/utils.rs @@ -265,7 +265,20 @@ impl RegionFileCopier { #[cfg(test)] mod tests { + #[cfg(feature = "hdfs-object-store")] + use object_store::ObjectStore; + #[cfg(feature = "hdfs-object-store")] + use object_store::layers::HdfsCompatibilityLayer; + #[cfg(feature = "hdfs-object-store")] + use object_store::services::Fs; + use super::*; + #[cfg(feature = "hdfs-object-store")] + use crate::access_layer::AccessLayer; + #[cfg(feature = "hdfs-object-store")] + use crate::sst::index::intermediate::IntermediateManager; + #[cfg(feature = "hdfs-object-store")] + use crate::sst::index::puffin_manager::PuffinManagerFactory; #[test] fn test_build_copy_file_paths() { @@ -344,4 +357,54 @@ mod tests { format!("/table_dir/1_0000000002/index/{}.1.puffin", file_id) ); } + + #[cfg(feature = "hdfs-object-store")] + #[tokio::test] + async fn test_copy_region_files_with_hdfs_fallback() { + let (temp_dir, puffin_manager) = + PuffinManagerFactory::new_for_test_async("hdfs-copy-region").await; + let intermediate_manager = IntermediateManager::init_fs(temp_dir.path().to_string_lossy()) + .await + .unwrap(); + let storage_dir = temp_dir.path().join("storage"); + std::fs::create_dir(&storage_dir).unwrap(); + let object_store = ObjectStore::new(Fs::default().root(storage_dir.to_str().unwrap())) + .unwrap() + .layer(HdfsCompatibilityLayer::new_for_test()); + let access_layer = Arc::new(AccessLayer::new( + "table_dir", + PathType::Bare, + object_store.clone(), + puffin_manager, + intermediate_manager, + )); + let copier = RegionFileCopier::new(access_layer); + let source_region_id = RegionId::new(1, 1); + let target_region_id = RegionId::new(1, 2); + let file_id = FileId::random(); + let descriptor = FileDescriptor::Data { file_id, size: 8 }; + let (source_path, target_path) = build_copy_file_paths( + source_region_id, + target_region_id, + descriptor, + "table_dir", + PathType::Bare, + ); + object_store.write(&source_path, "contents").await.unwrap(); + + copier + .copy_files(source_region_id, target_region_id, vec![descriptor], 1) + .await + .unwrap(); + + assert_eq!( + b"contents", + object_store + .read(&target_path) + .await + .unwrap() + .to_bytes() + .as_ref() + ); + } } diff --git a/src/mito2/src/region/version.rs b/src/mito2/src/region/version.rs index a56cef1f544..1c595fb0b9f 100644 --- a/src/mito2/src/region/version.rs +++ b/src/mito2/src/region/version.rs @@ -32,7 +32,7 @@ use store_api::storage::{RegionId, SequenceNumber}; use crate::error::Result; use crate::manifest::action::{RegionEdit, TruncateKind}; -use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef}; +use crate::memtable::time_partition::TimePartitions; use crate::memtable::version::{MemtableVersion, MemtableVersionRef}; use crate::memtable::{MemtableBuilderRef, MemtableId}; use crate::region::options::RegionOptions; @@ -319,13 +319,13 @@ impl VersionControl { version: &Version, metadata: RegionMetadataRef, memtable_builder: MemtableBuilderRef, - ) -> TimePartitionsRef { - Arc::new(TimePartitions::new( + ) -> TimePartitions { + TimePartitions::new( metadata, memtable_builder, version.memtables.mutable.next_memtable_id(), Some(version.memtables.mutable.part_duration()), - )) + ) } } @@ -402,7 +402,7 @@ pub(crate) struct VersionBuilder { impl VersionBuilder { /// Returns a new builder. - pub(crate) fn new(metadata: RegionMetadataRef, mutable: TimePartitionsRef) -> Self { + pub(crate) fn new(metadata: RegionMetadataRef, mutable: TimePartitions) -> Self { VersionBuilder { metadata: metadata.clone(), memtables: Arc::new(MemtableVersion::new(mutable)), diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index a64c37dac7f..f8d7eb9c840 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -70,8 +70,19 @@ impl Drop for WriteNotify { } } +/// Source of writes installed into memtables. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WriteSource { + Request, + WalReplay, +} + /// Context to keep region metadata and buffer write requests. pub(crate) struct RegionWriteCtx { + /// Source of the buffered writes. + source: WriteSource, + /// Whether accepted non-empty writes bypass WAL. WAL replay never sets this flag. + has_skip_wal_writes: bool, /// Id of region to write. region_id: RegionId, /// Version of the region while creating the context. @@ -120,6 +131,7 @@ impl RegionWriteCtx { version_control: &VersionControlRef, provider: Provider, written_bytes: Option>, + source: WriteSource, ) -> RegionWriteCtx { let VersionControlData { version, @@ -129,6 +141,8 @@ impl RegionWriteCtx { } = version_control.current(); RegionWriteCtx { + source, + has_skip_wal_writes: false, region_id, version, version_control: version_control.clone(), @@ -162,6 +176,7 @@ impl RegionWriteCtx { self.next_sequence = sequence; } let num_rows = rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0); + self.has_skip_wal_writes |= num_rows > 0 && self.is_skip_wal_write(skip_wal); let mutation = Mutation { op_type, sequence: self.next_sequence, @@ -255,8 +270,29 @@ impl RegionWriteCtx { self.next_entry_id } - /// Consumes mutations and writes them into mutable memtable. - pub(crate) async fn write_memtable(&mut self) { + /// Installs ordinary and bulk writes, then publishes their sequence and entry ID. + pub(crate) async fn write_memtables(&mut self) { + if self.failed { + return; + } + // A failed installation may leave some rows in memory. + if self.has_skip_wal_writes { + self.version.memtables.mutable.mark_skip_wal_writes(); + } + self.write_mutations().await; + self.write_bulk().await; + self.publish_sequence_and_entry_id(); + } + + fn is_skip_wal_write(&self, request_skip_wal: bool) -> bool { + self.source == WriteSource::Request + && (request_skip_wal + || self.provider == Provider::Noop + || self.version.options.skip_wal) + } + + /// Consumes mutations and writes them into mutable memtables. + async fn write_mutations(&mut self) { debug_assert_eq!(self.wal_notifiers.len(), self.wal_entry.mutations.len()); if self.failed { @@ -336,6 +372,7 @@ impl RegionWriteCtx { self.wal_entry.bulk_entries.push(entry); } + self.has_skip_wal_writes |= bulk.num_rows() > 0 && self.is_skip_wal_write(skip_wal); self.bulk_notifiers .push(WriteNotify::new(sender, bulk.num_rows())); @@ -344,7 +381,7 @@ impl RegionWriteCtx { true } - pub(crate) async fn write_bulk(&mut self) { + async fn write_bulk(&mut self) { if self.failed || self.bulk_parts.is_empty() { return; } @@ -398,10 +435,10 @@ impl RegionWriteCtx { } /// Publishes the assigned sequences and entry id to the region's committed - /// watermark. Only call after both [`write_memtable`](Self::write_memtable) + /// watermark. Only call after both [`write_mutations`](Self::write_mutations) /// and [`write_bulk`](Self::write_bulk) have completed; a failed context /// must not publish. - pub(crate) fn publish_sequence_and_entry_id(&self) { + fn publish_sequence_and_entry_id(&self) { if self.failed { return; } @@ -576,6 +613,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); for (op_type, skip, num_rows) in [ (OpType::Put, skip_wal, 2), @@ -657,6 +695,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); ctx.push_mutation( OpType::Delete as i32, @@ -693,6 +732,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); ctx.push_mutation( @@ -736,6 +776,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); // Alternate policies in one context, retaining all parts for the memtable. for skip in [skip_wal, false, skip_wal] { @@ -772,8 +813,13 @@ mod tests { 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::noop_provider(), None); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::noop_provider(), + None, + WriteSource::Request, + ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(), None, skip_wal)); @@ -791,6 +837,142 @@ mod tests { assert!(result.is_err(), "bulk notifier should report WAL error"); } + #[tokio::test] + async fn test_skip_wal_writes_source_and_policy() { + use api::v1::helper::{tag_column_schema, time_index_column_schema}; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, Value}; + + for source in [WriteSource::Request, WriteSource::WalReplay] { + for skip_wal in [false, true] { + for bulk in [false, true] { + for fail_before_install in [false, true] { + 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, + source, + ); + let (tx, rx) = oneshot::channel(); + if bulk { + let mut part = new_bulk_part(); + part.fill_missing_columns(&ctx.version.metadata).unwrap(); + assert!(ctx.push_bulk(OptionOutputTx::from(tx), part, None, skip_wal)); + } else { + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![ + time_index_column_schema( + "ts", + ColumnDataType::TimestampMillisecond, + ), + tag_column_schema("tag_0", ColumnDataType::String), + ], + rows: vec![Row { + values: vec![ + crate::test_util::ts_ms_value(0), + Value { + value_data: Some(ValueData::StringValue( + "a".to_string(), + )), + }, + ], + }], + }), + None, + OptionOutputTx::from(tx), + None, + skip_wal, + ); + } + if fail_before_install { + ctx.set_error(Arc::new( + UnexpectedSnafu { + reason: "WAL append failed".to_string(), + } + .build(), + )); + } + ctx.write_memtables().await; + assert_eq!( + version_control + .current() + .version + .memtables + .has_skip_wal_writes(), + source == WriteSource::Request && skip_wal && !fail_before_install, + "source={source:?}, skip_wal={skip_wal}, bulk={bulk}, failed={fail_before_install}", + ); + assert_eq!( + version_control.committed_sequence(), + if fail_before_install { + 0 + } else if bulk { + 2 + } else { + 1 + } + ); + drop(ctx); + assert_eq!(rx.await.unwrap().is_err(), fail_before_install); + } + } + } + } + } + + #[test] + fn test_skip_wal_writes_region_policy_and_empty_requests() { + for source in [WriteSource::Request, WriteSource::WalReplay] { + for noop in [false, true] { + for region_skip_wal in [false, true] { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let mut version = builder.build_version(); + version.options.skip_wal = region_skip_wal; + let version_control = + Arc::new(crate::region::version::VersionControl::new(version)); + let provider = if noop { + Provider::Noop + } else { + Provider::raft_engine_provider(region_id.as_u64()) + }; + let mut ctx = + RegionWriteCtx::new(region_id, &version_control, provider, None, source); + ctx.push_mutation( + OpType::Put as i32, + None, + None, + OptionOutputTx::none(), + None, + true, + ); + assert!(!ctx.has_skip_wal_writes); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default()], + }), + None, + OptionOutputTx::none(), + None, + false, + ); + assert_eq!( + ctx.has_skip_wal_writes, + source == WriteSource::Request && (noop || region_skip_wal) + ); + } + } + } + } + fn new_bulk_part() -> BulkPart { let schema = Arc::new(Schema::new(vec![Field::new( "ts", diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index 4a4133e1a9b..3aa70256709 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -39,7 +39,7 @@ use store_api::metric_engine_consts::{ use crate::error::Result; #[cfg(test)] -pub(crate) use crate::series_index::catalog::SeriesIndexEntry; +pub(crate) use crate::series_index::catalog::{RangeIndexEntry, SeriesIndexEntry}; pub(crate) use crate::series_index::catalog::{ delete_catalogs, load_version_control, series_index_path, }; diff --git a/src/mito2/src/series_index/bucket.rs b/src/mito2/src/series_index/bucket.rs index 7fecb253441..446948838f6 100644 --- a/src/mito2/src/series_index/bucket.rs +++ b/src/mito2/src/series_index/bucket.rs @@ -334,6 +334,7 @@ impl SeriesBucket { .filter_map(|file| file.meta_ref().sequence.map(|sequence| sequence.get())) .min()?; Some(SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: self.start, bucket_end: self.end, diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs index 8d361be364e..63a51badaeb 100644 --- a/src/mito2/src/series_index/builder.rs +++ b/src/mito2/src/series_index/builder.rs @@ -21,7 +21,6 @@ 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; @@ -34,7 +33,7 @@ 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, + RangeIndexEntry, SeriesIndexEntry, range_index_path, series_index_path, series_metadata, }; use crate::series_index::purger::{IndexFilePurger, IndexFileType, file_operation}; use crate::series_index::version::SeriesIndexFileHandle; @@ -68,7 +67,7 @@ pub(crate) async fn build_range_index( region: &MitoRegionRef, version: &VersionRef, file: FileHandle, -) -> Result> { +) -> Result> { let file_id = file.file_id().file_id(); let Some((context, mut selection)) = reader_input(region, file).await? else { return Ok(None); @@ -116,9 +115,12 @@ pub(crate) async fn build_range_index( } return Err(error); } - writer.finish().await?; + let metrics = writer.finish().await?; file_operation(IndexFileType::Range, "build", "success"); - Ok(Some(file_id)) + Ok(Some(RangeIndexEntry { + file_id, + file_size: metrics.output_bytes, + })) } /// Builds only the series index. Callers build needed range indexes separately. @@ -200,11 +202,14 @@ pub(crate) async fn build_series_index( } return Err(error); } - writer.finish().await?; + let metrics = writer.finish().await?; file_operation(IndexFileType::Series, "build", "success"); Ok(SeriesIndexFileHandle::new( region.region_id, - entry.clone(), + SeriesIndexEntry { + file_size: metrics.output_bytes, + ..entry.clone() + }, purger.clone(), )) } @@ -275,10 +280,19 @@ mod tests { let mut range_bytes = HashMap::new(); if build_ranges { for file in files { - let id = build_range_index(&store, ®ion, &version, file.clone()) + let completed = build_range_index(&store, ®ion, &version, file.clone()) .await .unwrap() .unwrap(); + let id = completed.file_id; + assert_eq!( + completed.file_size, + store + .stat(&range_index_path(region.region_id, id)) + .await + .unwrap() + .content_length() + ); range_bytes.insert( id, store @@ -314,7 +328,21 @@ mod tests { assert!(result.is_err()); } else { let handle = result.unwrap(); - assert_eq!(handle.entry(), &entry); + assert_eq!( + handle.entry().file_size, + store + .stat(&series_index_path(region.region_id, entry.index_uuid)) + .await + .unwrap() + .content_length() + ); + assert_eq!( + handle.entry(), + &SeriesIndexEntry { + file_size: handle.entry().file_size, + ..entry.clone() + } + ); assert!( store .exists(&series_index_path(region.region_id, entry.index_uuid)) @@ -504,12 +532,13 @@ mod tests { 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? + let Some(entry) = + 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); + completed.push(entry.file_id); } build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) .await diff --git a/src/mito2/src/series_index/catalog.rs b/src/mito2/src/series_index/catalog.rs index d56e10698fd..95cf87f83d3 100644 --- a/src/mito2/src/series_index/catalog.rs +++ b/src/mito2/src/series_index/catalog.rs @@ -60,6 +60,8 @@ pub(crate) struct WindowSequence { /// relies on this complete-coverage contract, not on `source_file_ids`. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct SeriesIndexEntry { + /// Completed size in the catalog; zero in the footer written before completion. + pub(crate) file_size: u64, pub(crate) index_uuid: FileId, /// Inclusive bucket start. pub(crate) bucket_start: Timestamp, @@ -97,6 +99,13 @@ impl SeriesIndexEntry { } } +/// A completed per-SST range index. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct RangeIndexEntry { + pub(crate) file_id: FileId, + pub(crate) file_size: u64, +} + #[derive(Debug, Default, Serialize, Deserialize)] pub(crate) struct SeriesIndexCatalog { pub(crate) indexes: Vec, @@ -104,7 +113,7 @@ pub(crate) struct SeriesIndexCatalog { #[derive(Debug, Default, Serialize, Deserialize)] pub(crate) struct RangeIndexCatalog { - pub(crate) indexes: Vec, + pub(crate) indexes: Vec, } pub(crate) fn range_catalog_path(region_id: RegionId) -> String { @@ -185,9 +194,12 @@ pub(crate) async fn load_version_control( let series = load_catalog::(store, &series_catalog_path(region_id)) .await .unwrap_or_default(); - // TODO: Handle catalog entries whose index files are missing from storage. let version = SeriesIndexVersion::new( - range.indexes.into_iter().collect(), + range + .indexes + .into_iter() + .map(|entry| (entry.file_id, entry)) + .collect(), series .indexes .into_iter() @@ -216,9 +228,9 @@ mod tests { use store_api::storage::{FileId, RegionId}; use crate::series_index::catalog::{ - RangeIndexCatalog, SeriesIndexCatalog, SeriesIndexEntry, WindowSequence, load_catalog, - load_version_control, range_catalog_path, series_catalog_path, series_metadata, - store_catalog, + RangeIndexCatalog, RangeIndexEntry, 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; @@ -250,6 +262,7 @@ mod tests { fn coverage_uses_exclusive_time_end_and_inclusive_file_sequences() { let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(1), bucket_end: Timestamp::new_second(2), @@ -304,7 +317,10 @@ mod tests { .write( &range_catalog_path(region_id), serde_json::to_vec(&RangeIndexCatalog { - indexes: vec![file_id], + indexes: vec![RangeIndexEntry { + file_id, + file_size: 1, + }], }) .unwrap(), ) @@ -315,7 +331,8 @@ mod tests { .await .unwrap(); let control = load_version_control(&store, region_id, &purger).await; - assert!(control.current().range_indexes.contains(&file_id)); + assert_eq!(1, control.current().range_indexes.len()); + assert_eq!(1, control.current().range_indexes[&file_id].file_size); assert!(control.current().series_indexes.is_empty()); let layer = MockLayerBuilder::default() .reader_factory(Arc::new(|_, _, _| Box::new(FailingCatalogReader))) @@ -342,6 +359,7 @@ mod tests { let store = ObjectStore::new(Memory::default()).unwrap(); let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(0), bucket_end: Timestamp::new_second(100), diff --git a/src/mito2/src/series_index/maintenance.rs b/src/mito2/src/series_index/maintenance.rs index d782cf5a2c1..f5221bb1ebe 100644 --- a/src/mito2/src/series_index/maintenance.rs +++ b/src/mito2/src/series_index/maintenance.rs @@ -20,7 +20,7 @@ use std::time::{Duration, Instant}; use common_telemetry::{debug, info}; use object_store::ObjectStore; -use store_api::storage::RegionId; +use store_api::storage::{FileId, RegionId}; use crate::error::Result; use crate::metrics::{SERIES_INDEX_RECONCILE_ELAPSED, SERIES_INDEX_RECONCILE_TOTAL}; @@ -74,6 +74,7 @@ impl ReconcileStats { } /// Reconciles indexes for one region snapshot, persists catalogs, then atomically publishes it. +#[allow(clippy::too_many_arguments)] pub(crate) async fn reconcile_series_indexes( worker_id: u32, store: ObjectStore, @@ -82,6 +83,7 @@ pub(crate) async fn reconcile_series_indexes( now_ms: i64, purger: IndexFilePurger, enable_range_index: bool, + allow_builds: bool, ) -> Result { let total_start = Instant::now(); // Use this snapshot throughout reconciliation, even if the region version advances. @@ -104,6 +106,7 @@ pub(crate) async fn reconcile_series_indexes( &purger, &mut unpublished, enable_range_index, + allow_builds, ) .await?; SERIES_INDEX_RECONCILE_ELAPSED @@ -164,6 +167,7 @@ async fn build_index_version( purger: &IndexFilePurger, unpublished: &mut UnpublishedSeriesFiles, enable_range_index: bool, + allow_builds: bool, ) -> Result<(Option, ReconcileStats)> { let mut stats = ReconcileStats::default(); let files = version @@ -179,25 +183,23 @@ async fn build_index_version( .map(|file| file.file_id().file_id()) .collect::>(); let current = region.series_index_version(); - // The SST purger deletes companion range files after final handle release. Prune - // metadata here using the captured SST snapshot, independently of physical deletion; - // a later region-version change is picked up by the next reconciliation. - stats.removed_range = current.range_indexes.difference(&visible).count(); - let buckets = match version.compaction_time_window { - Some(window) => rounded_bucket_width(requested_bucket_width, window) + // Empty build inputs still let the planner expire established bucket coverage. + let buckets = match (allow_builds, version.compaction_time_window) { + (true, Some(window)) => rounded_bucket_width(requested_bucket_width, window) // Successful rounding guarantees the window fits in i64; subsecond windows // use the same one-second minimum as rounded_bucket_width. .map(|width| { group_files_into_series_buckets(&files, width, (window.as_secs() as i64).max(1)) }) .unwrap_or_default(), - None => { + (true, None) => { debug!( "Deferring series indexes without compaction window, worker: {worker_id}, region: {}", region.region_id ); Vec::new() } + (false, _) => Vec::new(), }; let plan = plan_series_indexes( buckets, @@ -207,33 +209,41 @@ async fn build_index_version( ); stats.computed_buckets = plan.computed_buckets; stats.skipped_buckets = plan.skipped_buckets; + let mut next = prune_index_version(¤t, &visible, &plan.expired_index_ids, &mut stats); + if !allow_builds { + if let Some(next) = &mut next { + next.index_buckets = plan.index_buckets; + } + return Ok((next, stats)); + } if plan.builds.is_empty() - && plan.expired_index_ids.is_empty() - && stats.removed_range == 0 + && next.is_none() && (!enable_range_index || current.range_indexes.len() == visible.len()) { return Ok((None, stats)); } - let mut range_indexes = current.range_indexes.clone(); - range_indexes.retain(|file_id| visible.contains(file_id)); - let mut series_indexes = current.series_indexes.clone(); - for id in plan - .expired_index_ids - .iter() - .chain(&plan.superseded_index_ids) - { + let SeriesIndexVersion { + mut range_indexes, + mut series_indexes, + .. + } = next.unwrap_or_else(|| SeriesIndexVersion { + range_indexes: current.range_indexes.clone(), + series_indexes: current.series_indexes.clone(), + index_buckets: Default::default(), + }); + let index_buckets = plan.index_buckets; + for id in &plan.superseded_index_ids { series_indexes.remove(id); } for (bucket, expected) in plan.builds { // Complete companion indexes independently so a failed series build preserves them. for file in &bucket.files { if enable_range_index - && !range_indexes.contains(&file.file_id().file_id()) - && let Some(file_id) = - build_range_index(store, region, version, file.clone()).await? + && !range_indexes.contains_key(&file.file_id().file_id()) + && let Some(entry) = build_range_index(store, region, version, file.clone()).await? { stats.built_range += 1; - range_indexes.insert(file_id); + range_indexes.insert(entry.file_id, entry); } } let series_handle = @@ -246,12 +256,12 @@ async fn build_index_version( if enable_range_index { for file in files { let file_id = file.file_id().file_id(); - if range_indexes.contains(&file_id) { + if range_indexes.contains_key(&file_id) { continue; } - if let Some(file_id) = build_range_index(store, region, version, file).await? { + if let Some(entry) = build_range_index(store, region, version, file).await? { stats.built_range += 1; - range_indexes.insert(file_id); + range_indexes.insert(entry.file_id, entry); } } } @@ -268,11 +278,43 @@ async fn build_index_version( let next = SeriesIndexVersion { range_indexes, series_indexes, - index_buckets: plan.index_buckets, + index_buckets, }; Ok((Some(next), stats)) } +/// Prunes obsolete metadata without building indexes or retiring published handles. +/// Returns `None` without cloning the current version when no entries need removal. +/// Physical range deletion belongs to the SST purger; series handles are retired only +/// after the cleaned catalogs and snapshot have been published successfully. +fn prune_index_version( + current: &SeriesIndexVersion, + visible: &HashSet, + expired_index_ids: &[FileId], + stats: &mut ReconcileStats, +) -> Option { + if current.range_indexes.keys().all(|id| visible.contains(id)) + && expired_index_ids + .iter() + .all(|id| !current.series_indexes.contains_key(id)) + { + return None; + } + let mut next = SeriesIndexVersion { + range_indexes: current.range_indexes.clone(), + series_indexes: current.series_indexes.clone(), + index_buckets: current.index_buckets.clone(), + }; + next.range_indexes + .retain(|file_id, _| visible.contains(file_id)); + for id in expired_index_ids { + next.series_indexes.remove(id); + } + stats.removed_range = current.range_indexes.len() - next.range_indexes.len(); + stats.removed_series = current.series_indexes.len() - next.series_indexes.len(); + Some(next) +} + /// Writes changed catalogs in a stable order; the two writes are not atomic together. async fn persist_index_catalogs( store: &ObjectStore, @@ -281,8 +323,9 @@ async fn persist_index_catalogs( stats: &ReconcileStats, ) -> Result<()> { if stats.built_range + stats.removed_range > 0 { - let mut range_entries = next.range_indexes.iter().copied().collect::>(); - range_entries.sort_unstable_by(|left, right| left.as_bytes().cmp(right.as_bytes())); + let mut range_entries = next.range_indexes.values().copied().collect::>(); + range_entries + .sort_unstable_by(|left, right| left.file_id.as_bytes().cmp(right.file_id.as_bytes())); store_catalog( store, &range_catalog_path(region_id), @@ -328,3 +371,77 @@ fn publish_index_version(region: &MitoRegionRef, next: Arc) } } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use common_time::Timestamp; + use object_store::services::Memory; + + use super::*; + use crate::series_index::catalog::{RangeIndexEntry, SeriesIndexEntry}; + use crate::series_index::purger::series_index_channel; + + #[rstest::rstest] + fn test_prune_index_version( + #[values(false, true)] remove_range: bool, + #[values(false, true)] remove_series: bool, + ) { + let range_id = FileId::random(); + let series_id = FileId::random(); + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, mut receiver) = series_index_channel(store); + let current = SeriesIndexVersion::new( + HashMap::from([( + range_id, + RangeIndexEntry { + file_id: range_id, + file_size: 10, + }, + )]), + HashMap::from([( + series_id, + SeriesIndexFileHandle::new( + RegionId::new(1, 1), + SeriesIndexEntry { + index_uuid: series_id, + file_size: 20, + bucket_start: Timestamp::new_second(0), + bucket_end: Timestamp::new_second(100), + source_file_ids: vec![range_id], + min_file_sequence: 1, + max_file_sequence: 1, + compaction_window_secs: 100, + window_sequences: Default::default(), + }, + purger, + ), + )]), + ); + let visible = if remove_range { + HashSet::new() + } else { + HashSet::from([range_id]) + }; + // Unknown and repeated IDs must not inflate removal counts. + let mut expired = vec![FileId::random()]; + if remove_series { + expired.extend([series_id, series_id]); + } + let mut stats = ReconcileStats::default(); + let next = prune_index_version(¤t, &visible, &expired, &mut stats); + assert_eq!(remove_range || remove_series, next.is_some()); + assert_eq!(usize::from(remove_range), stats.removed_range); + assert_eq!(usize::from(remove_series), stats.removed_series); + if let Some(next) = next { + assert_eq!(!remove_range, next.range_indexes.contains_key(&range_id)); + assert_eq!(!remove_series, next.series_indexes.contains_key(&series_id)); + } + assert_eq!(1, current.range_indexes.len()); + assert_eq!(1, current.series_indexes.len()); + drop(current); + // Pruning alone must never retire published files. + assert!(receiver.try_recv().is_err()); + } +} diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs index 4613edae913..0b9befe4b8a 100644 --- a/src/mito2/src/series_index/searcher.rs +++ b/src/mito2/src/series_index/searcher.rs @@ -312,6 +312,7 @@ mod tests { ) -> (SeriesIndexFileHandle, UnboundedReceiver) { let (purger, receiver) = series_index_channel(store.clone()); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: common_time::Timestamp::new_second(0), bucket_end: common_time::Timestamp::new_second(60), diff --git a/src/mito2/src/series_index/task.rs b/src/mito2/src/series_index/task.rs index 2afbb245014..6ce340827a1 100644 --- a/src/mito2/src/series_index/task.rs +++ b/src/mito2/src/series_index/task.rs @@ -15,35 +15,57 @@ //! Worker-owned background maintenance for series indexes. use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; use common_telemetry::{info, warn}; use object_store::ObjectStore; -use tokio::sync::Notify; -use tokio::sync::mpsc::UnboundedReceiver; +use snafu::ensure; +use tokio::sync::mpsc::{self, Receiver, Sender, UnboundedReceiver}; +use tokio::sync::oneshot; use tokio::task::JoinHandle; -use tokio::time::{Instant, MissedTickBehavior}; +use tokio::time::Instant; -use crate::metrics::SERIES_INDEX_RECONCILE_TOTAL; -use crate::region::{RegionLeaderState, RegionMapRef, RegionRoleState}; -use crate::series_index::maintenance::reconcile_series_indexes; +use crate::error::{InvalidRequestSnafu, RegionBusySnafu, Result, WorkerStoppedSnafu}; +use crate::metrics::{ + SERIES_INDEX_CAPACITY_DEFERRED, SERIES_INDEX_DISK_BYTES, SERIES_INDEX_RECONCILE_TOTAL, +}; +use crate::read::series_candidate::is_sparse_metric_metadata; +use crate::region::{MitoRegionRef, RegionLeaderState, RegionMapRef, RegionRoleState}; +use crate::series_index::maintenance::{ReconcileStats, reconcile_series_indexes}; use crate::series_index::purger::{IndexFilePurger, PurgeRequest, run_index_purge_task}; use crate::time_provider::TimeProviderRef; -/// Shared lifecycle state for a worker's series-index task. +/// Commands serialized with periodic series-index maintenance. +#[derive(Debug)] +pub(crate) enum SeriesIndexCommand { + Wake, + Reconcile { + region: MitoRegionRef, + sender: oneshot::Sender>, + }, + Stop, +} + +/// Shared lifecycle state and command sender for a worker's series-index task. #[derive(Debug)] pub(crate) struct SeriesIndexTaskState { + worker_id: u32, running: AtomicBool, - notify: Notify, + sender: Sender, } impl SeriesIndexTaskState { - pub(crate) fn new() -> Self { - Self { - running: AtomicBool::new(true), - notify: Notify::new(), - } + pub(crate) fn new(worker_id: u32, channel_size: usize) -> (Self, Receiver) { + let (sender, receiver) = mpsc::channel(channel_size); + ( + Self { + worker_id, + running: AtomicBool::new(true), + sender, + }, + receiver, + ) } pub(crate) fn is_running(&self) -> bool { @@ -51,17 +73,37 @@ impl SeriesIndexTaskState { } pub(crate) fn wake(&self) { - self.notify.notify_one(); + // Queued work already wakes the task. Periodic maintenance still runs even + // if the queued command only reconciles one region. + if self.is_running() && self.sender.capacity() == self.sender.max_capacity() { + let _ = self.sender.try_send(SeriesIndexCommand::Wake); + } + } + + /// Returns a receiver for build completion, or RegionBusy if the queue is full. + pub(crate) fn try_reconcile( + &self, + region: MitoRegionRef, + ) -> Result>> { + ensure!(self.is_running(), WorkerStoppedSnafu { id: self.worker_id }); + let region_id = region.region_id; + let (sender, receiver) = oneshot::channel(); + self.sender + .try_send(SeriesIndexCommand::Reconcile { region, sender }) + .map_err(|error| match error { + mpsc::error::TrySendError::Full(_) => RegionBusySnafu { region_id }.build(), + mpsc::error::TrySendError::Closed(_) => { + WorkerStoppedSnafu { id: self.worker_id }.build() + } + })?; + Ok(receiver) } pub(crate) fn stop(&self) { - self.running.store(false, Ordering::Release); - // Retain a permit if maintenance has not started waiting yet. - self.notify.notify_one(); - } - - pub(crate) async fn notified(&self) { - self.notify.notified().await; + if self.running.swap(false, Ordering::AcqRel) { + // A full queue already wakes the task, which checks the running flag. + let _ = self.sender.try_send(SeriesIndexCommand::Stop); + } } } @@ -72,9 +114,12 @@ pub(crate) fn spawn_series_index_tasks( store: ObjectStore, regions: RegionMapRef, state: Arc, + receiver: Receiver, bucket_width: Duration, purger: IndexFilePurger, purge_receiver: UnboundedReceiver, + disk_usage: Arc, + max_size: u64, interval: Duration, time_provider: TimeProviderRef, enable_range_index: bool, @@ -92,7 +137,11 @@ pub(crate) fn spawn_series_index_tasks( regions, bucket_width, purger, + disk_usage, + max_size, + reported_usage: 0, state, + receiver, interval, time_provider, enable_range_index, @@ -108,8 +157,12 @@ struct SeriesIndexTask { regions: RegionMapRef, bucket_width: Duration, purger: IndexFilePurger, + disk_usage: Arc, + max_size: u64, + reported_usage: u64, worker_id: u32, state: Arc, + receiver: Receiver, interval: Duration, time_provider: TimeProviderRef, enable_range_index: bool, @@ -121,23 +174,125 @@ impl SeriesIndexTask { let worker_id = self.worker_id; info!("Start series-index background task, worker: {worker_id}"); let interval = self.time_provider.wait_duration(self.interval); - let mut timer = tokio::time::interval_at(Instant::now() + interval, interval); - // Schedule future ticks from a late tick rather than the original cadence. - timer.set_missed_tick_behavior(MissedTickBehavior::Delay); + let mut deadline = Instant::now() + interval; while self.state.is_running() { - tokio::select! { - _ = self.state.notified() => {} - _ = timer.tick() => {} - } - if self.state.is_running() { + // Check before receiving so a busy channel cannot starve periodic maintenance. + if Instant::now() >= deadline { self.maintain().await; + deadline = Instant::now() + interval; + continue; + } + match tokio::time::timeout_at(deadline, self.receiver.recv()).await { + Err(_) => continue, + Ok(None | Some(SeriesIndexCommand::Stop)) => break, + Ok(Some(SeriesIndexCommand::Wake)) => { + if self.state.is_running() { + self.maintain().await; + deadline = Instant::now() + interval; + } + } + Ok(Some(SeriesIndexCommand::Reconcile { region, sender })) => { + let result = self.reconcile_manual(region).await; + let _ = sender.send(result); + } + } + } + self.state.stop(); + self.receiver.close(); + while let Some(command) = self.receiver.recv().await { + if let SeriesIndexCommand::Reconcile { sender, .. } = command { + let _ = sender.send(WorkerStoppedSnafu { id: worker_id }.fail()); } } info!("Stop series-index background task, worker: {worker_id}"); } + /// A manual request must still refer to the same writable region instance. + fn validate_manual_region(&self, region: &MitoRegionRef) -> Result<()> { + ensure!( + self.state.is_running(), + WorkerStoppedSnafu { id: self.worker_id } + ); + let current = self.regions.writable_non_staging_region(region.region_id)?; + ensure!( + Arc::ptr_eq(¤t, region), + InvalidRequestSnafu { + region_id: region.region_id, + reason: "region was replaced during series-index reconciliation", + } + ); + ensure!( + is_sparse_metric_metadata(®ion.version().metadata), + InvalidRequestSnafu { + region_id: region.region_id, + reason: "series indexes require sparse metric metadata", + } + ); + Ok(()) + } + + async fn reconcile_manual(&mut self, region: MitoRegionRef) -> Result { + self.validate_manual_region(®ion)?; + self.refresh_usage(); + let result = self.reconcile_region(region.clone()).await; + self.refresh_usage(); + let result = result?; + self.validate_manual_region(®ion)?; + Ok(result) + } + + async fn reconcile_region(&self, region: MitoRegionRef) -> Result { + // Full capacity defers builds, but cleanup must still reclaim published usage. + let allow_builds = self.disk_usage.load(Ordering::Relaxed) < self.max_size; + if !allow_builds { + SERIES_INDEX_CAPACITY_DEFERRED.inc(); + } + let result = reconcile_series_indexes( + self.worker_id, + self.store.clone(), + region, + self.bucket_width, + self.time_provider.current_time_millis(), + self.purger.clone(), + self.enable_range_index, + allow_builds, + ) + .await; + if result.is_err() { + SERIES_INDEX_RECONCILE_TOTAL + .with_label_values(&["failure"]) + .inc(); + } + result + } + + /// Each worker contributes only its open regions, refreshed at maintenance boundaries. + fn refresh_usage(&mut self) { + let usage = self + .regions + .list_regions() + .iter() + .map(|region| region.series_index_version().disk_usage()) + .sum(); + self.report_usage(usage); + } + + fn report_usage(&mut self, usage: u64) { + if usage >= self.reported_usage { + let delta = usage - self.reported_usage; + self.disk_usage.fetch_add(delta, Ordering::Relaxed); + SERIES_INDEX_DISK_BYTES.add(delta as i64); + } else { + let delta = self.reported_usage - usage; + self.disk_usage.fetch_sub(delta, Ordering::Relaxed); + SERIES_INDEX_DISK_BYTES.sub(delta as i64); + } + self.reported_usage = usage; + } + /// Runs periodic maintenance independently of incoming deletion requests. async fn maintain(&mut self) { + self.refresh_usage(); for region in self.regions.list_regions() { if !self.state.is_running() { break; @@ -150,30 +305,26 @@ impl SeriesIndexTask { ) { continue; } - if let Err(error) = reconcile_series_indexes( - self.worker_id, - self.store.clone(), - region.clone(), - self.bucket_width, - self.time_provider.current_time_millis(), - self.purger.clone(), - self.enable_range_index, - ) - .await - { - SERIES_INDEX_RECONCILE_TOTAL - .with_label_values(&["failure"]) - .inc(); + if let Err(error) = self.reconcile_region(region.clone()).await { warn!(error; "Failed to reconcile series indexes, worker: {}, region: {}", self.worker_id, region.region_id); } + self.refresh_usage(); } } } +impl Drop for SeriesIndexTask { + fn drop(&mut self) { + self.report_usage(0); + } +} + #[cfg(test)] mod tests { + use object_store::layers::mock::{self, MockLayerBuilder, oio}; use object_store::services::Memory; - use store_api::region_engine::RegionEngine; + use store_api::region_engine::{RegionEngine, RegionRole}; + use tokio::sync::Notify; use super::*; use crate::region::RegionMap; @@ -182,6 +333,331 @@ mod tests { use crate::series_index::tests::prepare_region; use crate::test_util::TestEnv; + struct BlockingCatalogWriter { + inner: oio::Writer, + entered: Arc, + release: Arc, + } + + impl mock::Write for BlockingCatalogWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.inner.write(buffer).await + } + + async fn close(&mut self) -> mock::Result { + self.entered.notify_one(); + self.release.notified().await; + self.inner.close().await + } + + async fn abort(&mut self) -> mock::Result<()> { + self.inner.abort().await + } + } + + #[derive(Clone, Copy, Debug)] + enum LifecycleChange { + Remove, + Demote, + Stop, + } + + fn new_task(region: MitoRegionRef) -> (SeriesIndexTask, UnboundedReceiver) { + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, purge_receiver) = series_index_channel(store.clone()); + let regions = Arc::new(RegionMap::default()); + regions.insert_region(region); + let (state, receiver) = SeriesIndexTaskState::new(0, 1); + ( + SeriesIndexTask { + store, + regions, + bucket_width: Duration::from_secs(100), + purger, + disk_usage: Arc::default(), + max_size: u64::MAX, + reported_usage: 0, + worker_id: 0, + state: Arc::new(state), + receiver, + interval: Duration::from_secs(3600), + time_provider: Arc::new(crate::time_provider::StdTimeProvider), + enable_range_index: true, + }, + purge_receiver, + ) + } + + #[test] + fn test_coalesce_wakeups_and_stop() { + let (state, mut receiver) = SeriesIndexTaskState::new(0, 2); + state.wake(); + state.wake(); + assert!(matches!( + receiver.try_recv().unwrap(), + SeriesIndexCommand::Wake + )); + assert!(receiver.try_recv().is_err()); + state.wake(); + assert!(matches!( + receiver.try_recv().unwrap(), + SeriesIndexCommand::Wake + )); + state.stop(); + state.stop(); + state.wake(); + assert!(matches!( + receiver.try_recv().unwrap(), + SeriesIndexCommand::Stop + )); + assert!(receiver.try_recv().is_err()); + } + + #[rstest::rstest] + #[case(true)] + #[case(false)] + #[tokio::test] + async fn test_manual_reconcile_publishes_before_reply(#[case] range_enabled: bool) { + let mut env = TestEnv::with_prefix("series-manual").await; + let (engine, region) = prepare_region(&mut env).await; + let (mut task, _purge_receiver) = new_task(region.clone()); + task.enable_range_index = range_enabled; + let state = task.state.clone(); + let store = task.store.clone(); + let first = state.try_reconcile(region.clone()).unwrap(); + state.wake(); + assert_eq!(1, task.receiver.len()); + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::RegionBusy { .. }) + )); + let handle = tokio::spawn(task.run()); + let first = tokio::time::timeout(Duration::from_secs(10), first) + .await + .unwrap() + .unwrap() + .unwrap(); + // Admission recovers after completion; both builds reuse the same publication. + let second = state.try_reconcile(region.clone()).unwrap(); + let second = tokio::time::timeout(Duration::from_secs(10), second) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(1, first.built_series); + assert_eq!(if range_enabled { 4 } else { 0 }, first.built_range); + assert_eq!(0, second.built_series); + assert_eq!(0, second.built_range); + assert_eq!(1, region.series_index_version().series_indexes.len()); + assert!( + store + .exists(&series_catalog_path(region.region_id)) + .await + .unwrap() + ); + assert_eq!( + range_enabled, + store + .exists(&range_catalog_path(region.region_id)) + .await + .unwrap() + ); + state.stop(); + handle.await.unwrap(); + assert!(state.try_reconcile(region).is_err()); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_manual_reconcile_respects_capacity_and_refreshes_usage() { + let mut env = TestEnv::with_prefix("series-manual-capacity").await; + let (engine, region) = prepare_region(&mut env).await; + let (mut task, _purge_receiver) = new_task(region.clone()); + task.max_size = 0; + task.reconcile_manual(region.clone()).await.unwrap(); + assert_eq!(0, region.series_index_version().disk_usage()); + assert_eq!(0, task.disk_usage.load(Ordering::Relaxed)); + + task.max_size = u64::MAX; + task.reconcile_manual(region.clone()).await.unwrap(); + let bytes = region.series_index_version().disk_usage(); + assert!(bytes > 0); + assert_eq!(bytes, task.disk_usage.load(Ordering::Relaxed)); + let usage = task.disk_usage.clone(); + drop(task); + assert_eq!(0, usage.load(Ordering::Relaxed)); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_shutdown_fails_queued_reconciliation() { + let mut env = TestEnv::with_prefix("series-manual-stop").await; + let (engine, region) = prepare_region(&mut env).await; + let (task, _purge_receiver) = new_task(region.clone()); + let state = task.state.clone(); + let result = state.try_reconcile(region.clone()).unwrap(); + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::RegionBusy { .. }) + )); + state.stop(); + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::WorkerStopped { .. }) + )); + tokio::time::timeout(Duration::from_secs(10), async { + task.run().await; + assert!(matches!( + result.await.unwrap(), + Err(crate::error::Error::WorkerStopped { .. }) + )); + }) + .await + .unwrap(); + assert!(region.series_index_version().series_indexes.is_empty()); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_manual_reconcile_rejects_closed_region() { + let mut env = TestEnv::with_prefix("series-manual-closed").await; + let (engine, region) = prepare_region(&mut env).await; + let (task, _purge_receiver) = new_task(region.clone()); + let result = task.state.try_reconcile(region.clone()).unwrap(); + task.regions.remove_region(region.region_id); + let state = task.state.clone(); + let handle = tokio::spawn(task.run()); + assert!(matches!( + result.await.unwrap(), + Err(crate::error::Error::RegionNotFound { .. }) + )); + assert!(region.series_index_version().series_indexes.is_empty()); + state.stop(); + handle.await.unwrap(); + engine.stop().await.unwrap(); + } + + #[rstest::rstest] + #[case::remove(LifecycleChange::Remove)] + #[case::demote(LifecycleChange::Demote)] + #[case::stop(LifecycleChange::Stop)] + #[tokio::test] + async fn test_manual_reconcile_rejects_lifecycle_change_during_io( + #[case] change: LifecycleChange, + ) { + let mut env = TestEnv::with_prefix("series-manual-lifecycle").await; + let (engine, region) = prepare_region(&mut env).await; + let (mut task, _purge_receiver) = new_task(region.clone()); + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let writer_entered = entered.clone(); + let writer_release = release.clone(); + let catalog_path = series_catalog_path(region.region_id); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + if path == catalog_path { + Box::new(BlockingCatalogWriter { + inner, + entered: writer_entered.clone(), + release: writer_release.clone(), + }) + } else { + inner + } + })) + .build() + .unwrap(); + task.store = task.store.clone().layer(layer); + let state = task.state.clone(); + let regions = task.regions.clone(); + let mut result = state.try_reconcile(region.clone()).unwrap(); + let handle = tokio::spawn(task.run()); + tokio::time::timeout(Duration::from_secs(10), entered.notified()) + .await + .unwrap(); + assert!(matches!( + result.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + + // A canceled queued caller still occupies capacity while catalog I/O is blocked. + drop(state.try_reconcile(region.clone()).unwrap()); + for _ in 0..3 { + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::RegionBusy { .. }) + )); + } + + // The initial validation passed; invalidate the request during catalog I/O. + match change { + LifecycleChange::Remove => { + regions.remove_region(region.region_id).unwrap(); + } + LifecycleChange::Demote => region.set_role(RegionRole::Follower), + LifecycleChange::Stop => state.stop(), + } + release.notify_one(); + let result = tokio::time::timeout(Duration::from_secs(10), result) + .await + .unwrap() + .unwrap(); + state.stop(); + tokio::time::timeout(Duration::from_secs(10), handle) + .await + .unwrap() + .unwrap(); + engine.stop().await.unwrap(); + + match change { + LifecycleChange::Remove => assert!(matches!( + result, + Err(crate::error::Error::RegionNotFound { .. }) + )), + LifecycleChange::Demote => assert!(matches!( + result, + Err(crate::error::Error::RegionState { .. }) + )), + LifecycleChange::Stop => assert!(matches!( + result, + Err(crate::error::Error::WorkerStopped { .. }) + )), + } + } + + #[tokio::test] + async fn test_manual_traffic_does_not_postpone_periodic_maintenance() { + let mut env = TestEnv::with_prefix("series-manual-deadline").await; + let (engine, region) = prepare_region(&mut env).await; + region.switch_state_to_staging(RegionLeaderState::Writable); + region + .manifest_ctx + .exit_staging(region.region_id, RegionRoleState::Follower) + .unwrap(); + let (mut task, _purge_receiver) = new_task(region.clone()); + task.interval = Duration::from_millis(20); + let state = task.state.clone(); + let handle = tokio::spawn(task.run()); + // Manual builds on followers fail. Only a periodic sweep can publish coverage. + tokio::time::timeout(Duration::from_secs(10), async { + while region.series_index_version().series_indexes.is_empty() { + assert!( + state + .try_reconcile(region.clone()) + .unwrap() + .await + .unwrap() + .is_err() + ); + } + }) + .await + .unwrap(); + state.stop(); + handle.await.unwrap(); + engine.stop().await.unwrap(); + } + #[rstest::rstest] #[case::follower(RegionRoleState::Follower, true)] #[case::writable(RegionRoleState::Leader(RegionLeaderState::Writable), true)] @@ -206,17 +682,26 @@ mod tests { let (purger, _receiver) = series_index_channel(store.clone()); let regions = Arc::new(RegionMap::default()); regions.insert_region(region.clone()); + let (state, receiver) = SeriesIndexTaskState::new(0, 2); let mut task = SeriesIndexTask { store: store.clone(), regions, bucket_width: Duration::from_secs(100), purger, + disk_usage: Arc::default(), + max_size: u64::MAX, + reported_usage: 0, worker_id: 0, - state: Arc::new(SeriesIndexTaskState::new()), + state: Arc::new(state), + receiver, interval: Duration::from_secs(3600), time_provider: Arc::new(crate::time_provider::StdTimeProvider), enable_range_index: true, }; + assert_eq!( + role == RegionRoleState::Leader(RegionLeaderState::Writable), + task.validate_manual_region(®ion).is_ok(), + ); task.maintain().await; assert_eq!( builds, @@ -230,4 +715,242 @@ mod tests { } engine.stop().await.unwrap(); } + + #[rstest::rstest] + #[case::at_capacity(0, true)] + #[case::above_capacity(1, true)] + #[case::at_capacity_without_range(0, false)] + #[case::above_capacity_without_range(1, false)] + #[tokio::test] + async fn test_full_capacity_cleanup_and_recovery( + #[case] excess: u64, + #[case] enable_range_index: bool, + ) { + use std::sync::Mutex; + + use object_store::layers::mock::MockLayerBuilder; + + use crate::series_index::catalog::load_version_control; + use crate::series_index::tests::prepare_region_with_timestamps; + use crate::time_provider::mock::MockTimeProvider; + + let mut env = TestEnv::with_prefix("series-capacity-cleanup").await; + let (engine, region) = + prepare_region_with_timestamps(&mut env, &[1000, 2000, 3000, 4000, 5000]).await; + assert_eq!( + 5, + region + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .count() + ); + let mut options = region.version().options.clone(); + options.ttl = Some(common_time::TimeToLive::Duration(Duration::from_secs(100))); + region.version_control.alter_options(options); + let writes = Arc::new(Mutex::new(Vec::new())); + let captured = writes.clone(); + let store = ObjectStore::new(Memory::default()).unwrap().layer( + MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + captured.lock().unwrap().push(path.to_string()); + inner + })) + .build() + .unwrap(), + ); + let (purger, mut receiver) = series_index_channel(store.clone()); + let regions = Arc::new(RegionMap::default()); + regions.insert_region(region.clone()); + let clock = Arc::new(MockTimeProvider::new(0)); + let usage = Arc::new(AtomicU64::new(0)); + let (state, command_receiver) = SeriesIndexTaskState::new(0, 1); + let mut task = SeriesIndexTask { + store: store.clone(), + regions, + bucket_width: Duration::from_secs(100), + purger, + disk_usage: usage.clone(), + max_size: u64::MAX, + reported_usage: 0, + worker_id: 0, + state: Arc::new(state), + receiver: command_receiver, + interval: Duration::from_secs(3600), + time_provider: clock.clone(), + enable_range_index, + }; + task.maintain().await; + let previous = region.series_index_version(); + assert_eq!(1, previous.series_indexes.len()); + let old_id = *previous.series_indexes.keys().next().unwrap(); + task.max_size = previous.disk_usage() - excess; + writes.lock().unwrap().clear(); + task.maintain().await; + assert!(Arc::ptr_eq(&previous, ®ion.series_index_version())); + assert!(writes.lock().unwrap().is_empty()); + + // Removing the newest SST makes range coverage obsolete and would normally + // trigger a series replacement over the four remaining SSTs. + let sources = region.version(); + let newest = sources + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .max_by_key(|file| file.meta_ref().sequence) + .unwrap() + .meta_ref() + .clone(); + region.version_control.apply_edit( + Some(crate::manifest::action::RegionEdit { + files_to_remove: vec![newest.clone()], + files_to_add: Vec::new(), + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }), + &[], + crate::test_util::new_noop_file_purger(), + ); + task.maintain().await; + let cleaned = region.series_index_version(); + assert!(!cleaned.range_indexes.contains_key(&newest.file_id)); + assert_eq!( + usize::from(enable_range_index) * 4, + cleaned.range_indexes.len() + ); + assert_eq!(previous.index_buckets, cleaned.index_buckets); + assert_eq!(1, cleaned.series_indexes.len()); + assert!(cleaned.series_indexes.contains_key(&old_id)); + assert_eq!(cleaned.disk_usage(), usage.load(Ordering::Relaxed)); + assert_eq!( + previous.disk_usage() + - previous + .range_indexes + .get(&newest.file_id) + .map_or(0, |e| e.file_size), + cleaned.disk_usage() + ); + assert!( + writes + .lock() + .unwrap() + .iter() + .all(|path| path == &range_catalog_path(region.region_id)) + ); + let restored = load_version_control(&store, region.region_id, &task.purger).await; + assert_eq!(cleaned.range_indexes, restored.current().range_indexes); + assert_eq!(cleaned.index_buckets, restored.current().index_buckets); + drop(restored); + + // Expiration must work even when the shared estimate remains full. + task.max_size = cleaned.disk_usage() - excess; + clock.set_now(201_000); + writes.lock().unwrap().clear(); + task.maintain().await; + let expired = region.series_index_version(); + assert!(expired.series_indexes.is_empty()); + assert!(expired.index_buckets.is_empty()); + assert_eq!(cleaned.range_indexes, expired.range_indexes); + assert_eq!(expired.disk_usage(), usage.load(Ordering::Relaxed)); + assert!(usage.load(Ordering::Relaxed) < task.max_size); + assert_eq!( + *writes.lock().unwrap(), + vec![series_catalog_path(region.region_id)] + ); + let restored = load_version_control(&store, region.region_id, &task.purger).await; + assert!(restored.current().series_indexes.is_empty()); + assert_eq!(expired.range_indexes, restored.current().range_indexes); + assert!(receiver.try_recv().is_err()); + drop(previous); + assert!(receiver.try_recv().is_err()); + drop(cleaned); + assert_eq!(old_id, receiver.try_recv().unwrap().file_id.file_id()); + assert!(receiver.try_recv().is_err()); + + // Make the remaining SSTs eligible again; reclaimed capacity admits a build + // without closing the region or increasing the configured limit. + clock.set_now(0); + task.maintain().await; + let rebuilt = region.series_index_version(); + assert_eq!(1, rebuilt.series_indexes.len()); + assert!(!rebuilt.series_indexes.contains_key(&old_id)); + assert_eq!( + 4, + rebuilt + .series_indexes + .values() + .next() + .unwrap() + .entry() + .source_file_ids + .len() + ); + assert_eq!(rebuilt.disk_usage(), usage.load(Ordering::Relaxed)); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_shared_usage_defers_builds_and_allows_overshoot() { + let mut env = TestEnv::with_prefix("series-approximate-usage").await; + let (engine, region) = prepare_region(&mut env).await; + let usage = Arc::new(AtomicU64::new(0)); + let store = ObjectStore::new(Memory::default()).unwrap(); + let make_task = || { + let (state, command_receiver) = SeriesIndexTaskState::new(0, 1); + SeriesIndexTask { + store: store.clone(), + regions: Arc::new(RegionMap::default()), + bucket_width: Duration::from_secs(100), + purger: series_index_channel(store.clone()).0, + disk_usage: usage.clone(), + max_size: 1024, + reported_usage: 0, + worker_id: 0, + state: Arc::new(state), + receiver: command_receiver, + interval: Duration::from_secs(3600), + time_provider: Arc::new(crate::time_provider::StdTimeProvider), + enable_range_index: true, + } + }; + let mut task = make_task(); + task.regions.insert_region(region.clone()); + task.maintain().await; + let bytes = region.series_index_version().disk_usage(); + assert!( + bytes > task.max_size, + "a started reconciliation may exceed the limit" + ); + assert_eq!(bytes, usage.load(Ordering::Relaxed)); + // Another worker sees the same estimate and defers builds without creating coverage. + let mut other_env = TestEnv::with_prefix("series-approximate-other").await; + let (other_engine, other_region) = prepare_region(&mut other_env).await; + let mut other = make_task(); + other.regions.insert_region(other_region.clone()); + other.max_size = bytes; // Exact equality also skips. + other.maintain().await; + assert_eq!(0, other_region.series_index_version().disk_usage()); + assert_eq!(bytes, usage.load(Ordering::Relaxed)); + + // Closing a region reduces the estimate on the next pass, even while full. + task.regions.remove_region(region.region_id); + task.maintain().await; + assert_eq!(0, usage.load(Ordering::Relaxed)); + other.maintain().await; + assert!(other_region.series_index_version().disk_usage() > 0); + assert_eq!( + other_region.series_index_version().disk_usage(), + usage.load(Ordering::Relaxed) + ); + drop(other); + assert_eq!(0, usage.load(Ordering::Relaxed)); + engine.stop().await.unwrap(); + other_engine.stop().await.unwrap(); + } } diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs index c1efe3afdcf..631bdc75b87 100644 --- a/src/mito2/src/series_index/tests.rs +++ b/src/mito2/src/series_index/tests.rs @@ -48,11 +48,27 @@ pub(super) async fn prepare_region(env: &mut TestEnv) -> (MitoEngine, MitoRegion prepare_region_with_timestamps(env, &[1000, 2000, 3000, 4000]).await } -async fn prepare_region_with_timestamps( +pub(crate) async fn prepare_region_with_timestamps( env: &mut TestEnv, timestamps: &[i64], ) -> (MitoEngine, MitoRegionRef) { - let engine = env.create_engine(MitoConfig::default()).await; + prepare_region_with_config( + env, + timestamps, + MitoConfig { + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }, + ) + .await +} + +async fn prepare_region_with_config( + env: &mut TestEnv, + timestamps: &[i64], + config: MitoConfig, +) -> (MitoEngine, MitoRegionRef) { + let engine = env.create_engine(config).await; let metadata = Arc::new(sst_region_metadata_with_encoding( PrimaryKeyEncoding::Sparse, )); @@ -93,6 +109,10 @@ async fn prepare_region_with_timestamps( .handle_request(region_id, RegionRequest::Create(request)) .await .unwrap(); + // Flush completion can precede compaction scheduling. Keep background TTL + // cleanup from racing tests that alter options after preparing the SSTs. + let region = engine.get_region(region_id).unwrap(); + region.update_schedule_compaction_millis(); for &ts in timestamps { engine .handle_request( @@ -122,7 +142,6 @@ async fn prepare_region_with_timestamps( .unwrap(); flush_region(&engine, region_id, None).await; } - let region = engine.get_region(region_id).unwrap(); (engine, region) } @@ -171,6 +190,7 @@ impl IndexTest { 0, self.purger.clone(), enable_range_index, + true, ) .await } @@ -311,7 +331,7 @@ async fn test_reconcile_restores_and_reuses_indexes() { let missing_paths = [ range_index_path( region.region_id, - *first.range_indexes.iter().next().unwrap(), + *first.range_indexes.keys().next().unwrap(), ), series_index_path(region.region_id, first_id), ]; @@ -325,6 +345,7 @@ async fn test_reconcile_restores_and_reuses_indexes() { .current(); assert_eq!(first.range_indexes, restored.range_indexes); assert_eq!(first.index_buckets, restored.index_buckets); + assert_eq!(first.disk_usage(), restored.disk_usage()); assert_eq!( first.series_indexes[&first_id].entry(), restored.series_indexes[&first_id].entry() @@ -697,13 +718,13 @@ async fn test_failed_reconcile_retires_only_unpublished_series() { ); assert_eq!(if failure == "range" { 2 } else { 3 }, plan.builds.len()); let (bucket, entry) = &plan.builds[0]; - let mut ranges = HashSet::new(); + let mut ranges = HashMap::new(); for file in &bucket.files { - let id = build_range_index(store, ®ion, &version, file.clone()) + let entry = build_range_index(store, ®ion, &version, file.clone()) .await .unwrap() .unwrap(); - ranges.insert(id); + ranges.insert(entry.file_id, entry); } let handle = build_series_index(store, ®ion, &version, bucket, entry, purger) .await @@ -824,7 +845,8 @@ async fn test_maintenance_wakeup_and_timer(#[case] enable_range_index: bool) { region.version_control.alter_options(options); let store = ObjectStore::new(Memory::default()).unwrap(); let (purger, receiver) = series_index_channel(store.clone()); - let state = Arc::new(super::task::SeriesIndexTaskState::new()); + let (state, commands) = super::task::SeriesIndexTaskState::new(0, 2); + let state = Arc::new(state); let clock = Arc::new(crate::time_provider::mock::MockTimeProvider::new(0)); // A notification issued before the task starts must also trigger maintenance. state.wake(); @@ -835,9 +857,12 @@ async fn test_maintenance_wakeup_and_timer(#[case] enable_range_index: bool) { store, regions, state.clone(), + commands, Duration::from_secs(100), purger, receiver, + Arc::default(), + u64::MAX, Duration::from_secs(3600), clock.clone(), enable_range_index, @@ -866,6 +891,7 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { let store = ObjectStore::new(Memory::default()).unwrap(); let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(0), bucket_end: Timestamp::new_second(100), @@ -904,16 +930,20 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { let control = load_version_control(&store, region_id, &purger).await; let snapshot = control.current(); assert_eq!(&entry, snapshot.series_indexes[&entry.index_uuid].entry()); - let state = Arc::new(super::task::SeriesIndexTaskState::new()); + let (state, commands) = super::task::SeriesIndexTaskState::new(0, 2); + let state = Arc::new(state); state.stop(); super::task::spawn_series_index_tasks( 0, store.clone(), Arc::new(RegionMap::default()), state, + commands, Duration::from_secs(100), purger, receiver, + Arc::default(), + u64::MAX, Duration::from_secs(3600), Arc::new(StdTimeProvider), true, @@ -936,3 +966,113 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { .await .unwrap(); } + +#[rstest::rstest] +#[case(true)] +#[case(false)] +#[tokio::test] +async fn test_build_series_index_region_request(#[case] enabled: bool) { + use api::v1::region::build_index_request; + use store_api::region_request::RegionBuildIndexRequest; + + let mut env = TestEnv::with_prefix("series-manual-request").await; + let config = MitoConfig { + experimental_enable_series_index: enabled, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let (engine, region) = + prepare_region_with_config(&mut env, &[1000, 2000, 3000, 4000], config).await; + for _ in 0..2 { + let result = engine + .handle_request( + region.region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest { + options: Some(build_index_request::Options::SeriesIndex(Default::default())), + }), + ) + .await; + if enabled { + assert_eq!(0, result.unwrap().affected_rows); + assert!(!region.series_index_version().series_indexes.is_empty()); + } else { + assert!( + result + .unwrap_err() + .to_string() + .contains("series index is disabled") + ); + } + } + engine.stop().await.unwrap(); +} + +#[rstest::rstest] +#[case::build(false)] +#[case::catalog(true)] +#[tokio::test] +async fn test_manual_reconcile_returns_storage_failure(#[case] catalog_failure: bool) { + let mut env = TestEnv::with_prefix("series-manual-failure").await; + let (engine, region) = prepare_region(&mut env).await; + let catalog = series_catalog_path(region.region_id); + let fail_once = Arc::new(std::sync::atomic::AtomicBool::new(true)); + let fail = fail_once.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + let target = if catalog_failure { + path == catalog + } else { + path.contains("/series/") + }; + Box::new(FailingSeriesWriter { + inner, + fail: target && fail.load(std::sync::atomic::Ordering::Relaxed), + }) + })) + .build() + .unwrap(); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + let (purger, purge_receiver) = series_index_channel(store.clone()); + let (state, receiver) = super::task::SeriesIndexTaskState::new(0, 2); + let state = Arc::new(state); + let regions = Arc::new(RegionMap::default()); + regions.insert_region(region.clone()); + let task = super::task::spawn_series_index_tasks( + 0, + store, + regions, + state.clone(), + receiver, + Duration::from_secs(100), + purger, + purge_receiver, + Arc::default(), + u64::MAX, + Duration::from_secs(3600), + Arc::new(StdTimeProvider), + false, + ); + assert!( + state + .try_reconcile(region.clone()) + .unwrap() + .await + .unwrap() + .is_err() + ); + assert!(region.series_index_version().series_indexes.is_empty()); + fail_once.store(false, std::sync::atomic::Ordering::Relaxed); + assert_eq!( + 1, + state + .try_reconcile(region) + .unwrap() + .await + .unwrap() + .unwrap() + .built_series + ); + state.stop(); + task.await.unwrap(); + engine.stop().await.unwrap(); +} diff --git a/src/mito2/src/series_index/version.rs b/src/mito2/src/series_index/version.rs index c31dfde906b..fa2115fc1f2 100644 --- a/src/mito2/src/series_index/version.rs +++ b/src/mito2/src/series_index/version.rs @@ -14,7 +14,7 @@ //! Immutable index snapshots and aggregate series-file handles. -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap}; use std::fmt::{self, Debug, Formatter}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; @@ -23,7 +23,7 @@ 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::catalog::{RangeIndexEntry, SeriesIndexEntry}; use crate::series_index::purger::{IndexFilePurger, PurgeRequest}; use crate::sst::file::RegionFileId; @@ -94,7 +94,7 @@ impl Drop for SeriesIndexFileHandleInner { 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) range_indexes: HashMap, pub(crate) series_indexes: HashMap, pub(crate) index_buckets: BTreeMap, } @@ -102,7 +102,7 @@ pub(crate) struct SeriesIndexVersion { impl SeriesIndexVersion { /// Restores bucket lookup from immutable index coverage stored in the catalog. pub(crate) fn new( - range_indexes: HashSet, + range_indexes: HashMap, series_indexes: HashMap, ) -> Self { let mut index_buckets = BTreeMap::new(); @@ -116,6 +116,19 @@ impl SeriesIndexVersion { } } + /// Approximate installed usage; old snapshots and unpublished outputs are excluded. + pub(crate) fn disk_usage(&self) -> u64 { + self.range_indexes + .values() + .map(|entry| entry.file_size) + .sum::() + + self + .series_indexes + .values() + .map(|handle| handle.entry().file_size) + .sum::() + } + fn mark_all_deleted(&self) { self.series_indexes .values() diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index f160bbb3785..87684586fbf 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -298,20 +298,33 @@ fn is_false(value: &bool) -> bool { !*value } +/// Formats a debug field with a custom closure, as a stable replacement for +/// the unstable `DebugStruct::field_with`. +struct DebugFmt) -> fmt::Result>(F); + +impl) -> fmt::Result> std::fmt::Debug for DebugFmt { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + (self.0)(f) + } +} + impl Debug for FileMeta { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { let mut debug_struct = f.debug_struct("FileMeta"); debug_struct .field("region_id", &self.region_id) - .field_with("file_id", |f| write!(f, "{} ", self.file_id)) - .field_with("time_range", |f| { - write!( - f, - "({}, {}) ", - self.time_range.0.to_iso8601_string(), - self.time_range.1.to_iso8601_string() - ) - }) + .field("file_id", &DebugFmt(|f| write!(f, "{} ", self.file_id))) + .field( + "time_range", + &DebugFmt(|f| { + write!( + f, + "({}, {}) ", + self.time_range.0.to_iso8601_string(), + self.time_range.1.to_iso8601_string() + ) + }), + ) .field("level", &self.level) .field("file_size", &ReadableSize(self.file_size)) .field( @@ -327,14 +340,17 @@ impl Debug for FileMeta { debug_struct .field("num_rows", &self.num_rows) .field("num_row_groups", &self.num_row_groups) - .field_with("sequence", |f| match self.sequence { - None => { - write!(f, "None") - } - Some(seq) => { - write!(f, "{}", seq) - } - }) + .field( + "sequence", + &DebugFmt(|f| match self.sequence { + None => { + write!(f, "None") + } + Some(seq) => { + write!(f, "{}", seq) + } + }), + ) .field("partition_expr", &self.partition_expr) .field("num_series", &self.num_series); if self.primary_key_min.is_some() || self.primary_key_max.is_some() { @@ -361,8 +377,7 @@ pub enum IndexType { FulltextIndex, /// Bloom Filter index BloomFilterIndex, - /// Vector index (HNSW). - #[cfg(feature = "vector_index")] + /// Retained for manifests written with the removed experimental vector index. VectorIndex, } @@ -426,12 +441,6 @@ impl FileMeta { .contains(&IndexType::BloomFilterIndex) } - /// Returns true if the file has a vector index. - #[cfg(feature = "vector_index")] - pub fn vector_index_available(&self) -> bool { - self.available_indexes.contains(&IndexType::VectorIndex) - } - pub fn index_file_size(&self) -> u64 { self.index_file_size } @@ -1062,6 +1071,31 @@ mod tests { assert_eq!(file_meta, deserialized_file_meta); } + #[test] + fn test_deserialize_legacy_vector_index() { + let json = r#"{ + "region_id": 0, + "file_id": "bc5896ec-e4d8-4017-a80d-f2de73188d55", + "time_range": [{"value":0,"unit":"Millisecond"},{"value":0,"unit":"Millisecond"}], + "available_indexes": ["VectorIndex", "InvertedIndex"], + "indexes": [{"column_id": 1, "created_indexes": ["VectorIndex"]}], + "level": 0 + }"#; + let meta: FileMeta = serde_json::from_str(json).unwrap(); + assert!(meta.inverted_index_available()); + assert!(!meta.fulltext_index_available()); + assert!(!meta.bloom_filter_index_available()); + let encoded = serde_json::to_value(&meta).unwrap(); + assert_eq!( + encoded["available_indexes"], + serde_json::json!(["VectorIndex", "InvertedIndex"]) + ); + assert_eq!( + encoded["indexes"][0]["created_indexes"], + serde_json::json!(["VectorIndex"]) + ); + } + #[test] fn test_file_meta_with_partition_expr() { let file_id = FileId::random(); diff --git a/src/mito2/src/sst/file_ref.rs b/src/mito2/src/sst/file_ref.rs index 2f067cf3c19..28ff116cf25 100644 --- a/src/mito2/src/sst/file_ref.rs +++ b/src/mito2/src/sst/file_ref.rs @@ -77,12 +77,8 @@ impl FileReferenceManager { } fn ref_file_set(&self, region_id: RegionId) -> Option> { - let file_refs = if let Some(file_refs) = self.files_per_region.get(®ion_id) { - file_refs.clone() - } else { - // region id not found. - return None; - }; + // region id not found. + let file_refs = self.files_per_region.get(®ion_id)?.clone(); if file_refs.files.is_empty() { // still return an empty manifest to indicate no files are referenced. diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 0c3ede4fb06..be52170f713 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -26,8 +26,6 @@ pub mod puffin_manager; mod sparse_test; mod statistics; pub(crate) mod store; -#[cfg(feature = "vector_index")] -pub(crate) mod vector_index; use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; @@ -37,28 +35,26 @@ use std::sync::Arc; use bloom_filter::creator::BloomFilterIndexer; use common_telemetry::{debug, error, info, warn}; use datatypes::arrow::record_batch::RecordBatch; +use mito_codec::row_converter::DensePrimaryKeyCodec; use object_store::ObjectStore; use puffin_manager::SstPuffinManager; use smallvec::{SmallVec, smallvec}; use snafu::ResultExt; use statistics::{ByteCount, RowCount}; +use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId, RegionId}; use strum::IntoStaticStr; use tokio::sync::mpsc::Sender; -#[cfg(feature = "vector_index")] -use vector_index::creator::VectorIndexer; use crate::access_layer::{AccessLayerRef, FilePathProvider, OperationType, RegionFilePathFactory}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; use crate::cache::write_cache::{UploadOptions, UploadTracker, WriteCacheRef}; use crate::cache::{CacheManagerRef, CacheStrategy}; -#[cfg(feature = "vector_index")] -use crate::config::VectorIndexConfig; use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig}; use crate::error::{ - BuildIndexAsyncSnafu, Error, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, - Result, + BuildIndexAsyncSnafu, Error, JoinSnafu, RegionClosedSnafu, RegionDroppedSnafu, + RegionTruncatedSnafu, Result, }; use crate::metrics::{ INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_MEMORY_USAGE, INDEX_PUBLICATION_STALE_TOTAL, @@ -87,8 +83,6 @@ use crate::worker::WorkerListener; pub(crate) const TYPE_INVERTED_INDEX: &str = "inverted_index"; pub(crate) const TYPE_FULLTEXT_INDEX: &str = "fulltext_index"; pub(crate) const TYPE_BLOOM_FILTER_INDEX: &str = "bloom_filter_index"; -#[cfg(feature = "vector_index")] -pub(crate) const TYPE_VECTOR_INDEX: &str = "vector_index"; /// Evicts local state for a stale index artifact. /// @@ -168,9 +162,6 @@ pub struct IndexOutput { pub fulltext_index: FulltextIndexOutput, /// Bloom filter output. pub bloom_filter: BloomFilterOutput, - /// Vector index output. - #[cfg(feature = "vector_index")] - pub vector_index: VectorIndexOutput, } impl IndexOutput { @@ -185,10 +176,6 @@ impl IndexOutput { if self.bloom_filter.is_available() { indexes.push(IndexType::BloomFilterIndex); } - #[cfg(feature = "vector_index")] - if self.vector_index.is_available() { - indexes.push(IndexType::VectorIndex); - } indexes } @@ -212,12 +199,6 @@ impl IndexOutput { .push(IndexType::BloomFilterIndex); } } - #[cfg(feature = "vector_index")] - if self.vector_index.is_available() { - for &col in &self.vector_index.columns { - map.entry(col).or_default().push(IndexType::VectorIndex); - } - } map.into_iter() .map(|(column_id, created_indexes)| ColumnIndexMetadata { @@ -251,9 +232,6 @@ pub type InvertedIndexOutput = IndexBaseOutput; pub type FulltextIndexOutput = IndexBaseOutput; /// Output of the bloom filter creation. pub type BloomFilterOutput = IndexBaseOutput; -/// Output of the vector index creation. -#[cfg(feature = "vector_index")] -pub type VectorIndexOutput = IndexBaseOutput; /// The index creator that hides the error handling details. #[derive(Default)] @@ -272,14 +250,60 @@ pub struct Indexer { last_mem_fulltext_index: usize, bloom_filter_indexer: Option, last_mem_bloom_filter: usize, - #[cfg(feature = "vector_index")] - vector_indexer: Option, - #[cfg(feature = "vector_index")] - last_mem_vector_index: usize, + /// Present only when the active creators together need every Dense PK field. + dense_pk_decoder: Option, intermediate_manager: Option, } impl Indexer { + /// Wraps real creators for update-only benchmarks without a Puffin output. + #[cfg(feature = "testing")] + pub fn for_bench( + metadata: &RegionMetadataRef, + inverted_indexer: Option, + bloom_filter_indexer: Option, + ) -> Self { + let mut indexer = Self { + region_id: metadata.region_id, + inverted_indexer, + bloom_filter_indexer, + ..Default::default() + }; + indexer.prepare_dense_pk_decoder(metadata); + indexer + } + + /// Computes demand once from active creators. Field columns and duplicate + /// requests across creators cannot turn a partial PK projection into a full one. + fn prepare_dense_pk_decoder(&mut self, metadata: &RegionMetadataRef) { + if metadata.primary_key_encoding != PrimaryKeyEncoding::Dense + || metadata.primary_key.is_empty() + { + return; + } + let requested: HashSet<_> = self + .inverted_indexer + .iter() + .flat_map(|indexer| indexer.column_ids()) + .chain( + self.bloom_filter_indexer + .iter() + .flat_map(|indexer| indexer.column_ids()), + ) + .collect(); + if metadata.primary_key.iter().all(|id| requested.contains(id)) { + self.dense_pk_decoder = Some(DensePrimaryKeyCodec::new(metadata)); + } + } + + /// Called before any creator consumes the batch, so the full decode is shared. + fn prepare_primary_key(&self, batch: &mut Batch) -> Result<()> { + if let Some(codec) = &self.dense_pk_decoder { + batch.ensure_dense_pk_decoded(codec)?; + } + Ok(()) + } + /// Updates the index with the given batch. pub async fn update(&mut self, batch: &mut Batch) { self.do_update(batch).await; @@ -336,18 +360,6 @@ impl Indexer { .with_label_values(&[TYPE_BLOOM_FILTER_INDEX]) .add(bloom_filter_mem as i64 - self.last_mem_bloom_filter as i64); self.last_mem_bloom_filter = bloom_filter_mem; - - #[cfg(feature = "vector_index")] - { - let vector_mem = self - .vector_indexer - .as_ref() - .map_or(0, |creator| creator.memory_usage()); - INDEX_CREATE_MEMORY_USAGE - .with_label_values(&[TYPE_VECTOR_INDEX]) - .add(vector_mem as i64 - self.last_mem_vector_index as i64); - self.last_mem_vector_index = vector_mem; - } } } @@ -372,8 +384,6 @@ pub(crate) struct IndexerBuilderImpl { pub(crate) inverted_index_config: InvertedIndexConfig, pub(crate) fulltext_index_config: FulltextIndexConfig, pub(crate) bloom_filter_index_config: BloomFilterConfig, - #[cfg(feature = "vector_index")] - pub(crate) vector_index_config: VectorIndexConfig, } #[async_trait::async_trait] @@ -398,18 +408,9 @@ impl IndexerBuilder for IndexerBuilderImpl { self.build_inverted_indexer(region_file_id.file_id(), row_group_size); indexer.fulltext_indexer = self.build_fulltext_indexer(region_file_id.file_id()).await; indexer.bloom_filter_indexer = self.build_bloom_filter_indexer(region_file_id.file_id()); - #[cfg(feature = "vector_index")] - { - indexer.vector_indexer = self.build_vector_indexer(region_file_id.file_id()); - } + indexer.prepare_dense_pk_decoder(&self.metadata); indexer.intermediate_manager = Some(self.intermediate_manager.clone()); - #[cfg(feature = "vector_index")] - let has_any_indexer = indexer.inverted_indexer.is_some() - || indexer.fulltext_indexer.is_some() - || indexer.bloom_filter_indexer.is_some() - || indexer.vector_indexer.is_some(); - #[cfg(not(feature = "vector_index"))] let has_any_indexer = indexer.inverted_indexer.is_some() || indexer.fulltext_indexer.is_some() || indexer.bloom_filter_indexer.is_some(); @@ -588,69 +589,6 @@ impl IndexerBuilderImpl { None } - - #[cfg(feature = "vector_index")] - fn build_vector_indexer(&self, file_id: FileId) -> Option { - let create = match self.build_type { - IndexBuildType::Flush => self.vector_index_config.create_on_flush.auto(), - IndexBuildType::Compact => self.vector_index_config.create_on_compaction.auto(), - _ => true, - }; - - if !create { - debug!( - "Skip creating vector index due to config, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - return None; - } - - // Get vector index column IDs and options from metadata - let vector_index_options = self.metadata.vector_indexed_column_ids(); - if vector_index_options.is_empty() { - debug!( - "No vector columns to index, skip creating vector index, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - return None; - } - - let mem_limit = self.vector_index_config.mem_threshold_on_create(); - let indexer = VectorIndexer::new( - file_id, - &self.metadata, - self.intermediate_manager.clone(), - mem_limit, - &vector_index_options, - ); - - let err = match indexer { - Ok(indexer) => { - if indexer.is_none() { - debug!( - "Skip creating vector index due to no columns require indexing, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - } - return indexer; - } - Err(err) => err, - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to create vector index, region_id: {}, file_id: {}, err: {:?}", - self.metadata.region_id, file_id, err - ); - } else { - warn!( - err; "Failed to create vector index, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - } - - None - } } /// Type of an index build task. @@ -774,7 +712,21 @@ impl IndexBuildTask { self.source.file_meta.file_id, )) .await; - match self.index_build(version_control).await { + let result = if self.reason == IndexBuildType::Compact { + let mut task = self.clone(); + // Keep the scheduler slot occupied until the compact runtime finishes the build. + match common_runtime::spawn_compact( + async move { task.index_build(version_control).await }, + ) + .await + { + Ok(result) => result, + Err(err) => Err(err).context(JoinSnafu), + } + } else { + self.index_build(version_control).await + }; + match result { Ok(outcome) => self.on_success(outcome).await, Err(e) => { warn!( @@ -1498,8 +1450,6 @@ mod tests { with_inverted: bool, with_fulltext: bool, with_skipping_bloom: bool, - #[cfg(feature = "vector_index")] - with_vector: bool, } async fn seed_manifest_file(manifest_ctx: &ManifestContextRef, file_meta: &FileMeta) { @@ -1526,8 +1476,6 @@ mod tests { with_inverted, with_fulltext, with_skipping_bloom, - #[cfg(feature = "vector_index")] - with_vector, }: MetaConfig, ) -> RegionMetadataRef { let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2)); @@ -1593,24 +1541,6 @@ mod tests { builder.push_column_metadata(column); } - #[cfg(feature = "vector_index")] - if with_vector { - use index::vector::VectorIndexOptions; - - let options = VectorIndexOptions::default(); - let column_schema = - ColumnSchema::new("vec", ConcreteDataType::vector_datatype(4), true) - .with_vector_index_options(&options) - .unwrap(); - let column = ColumnMetadata { - column_schema, - semantic_type: SemanticType::Field, - column_id: 6, - }; - - builder.push_column_metadata(column); - } - Arc::new(builder.build().unwrap()) } @@ -1658,7 +1588,6 @@ mod tests { op_type: OperationType::Flush, metadata: metadata.clone(), source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: Default::default(), @@ -1668,8 +1597,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut metrics = Metrics::new(WriteType::Flush); env.access_layer @@ -1684,12 +1611,12 @@ mod tests { file_purger: FilePurgerRef, files: HashMap, ) -> VersionControlRef { - let mutable = Arc::new(TimePartitions::new( + let mutable = TimePartitions::new( metadata.clone(), Arc::new(EmptyMemtableBuilder::default()), 0, None, - )); + ); let version_builder = VersionBuilder::new(metadata, mutable) .add_files(file_purger, files.values().cloned()) .build(); @@ -1719,8 +1646,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }) } @@ -1734,8 +1659,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1747,8 +1670,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1768,8 +1689,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1784,8 +1703,6 @@ mod tests { }, fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1807,8 +1724,6 @@ mod tests { ..Default::default() }, bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1830,8 +1745,6 @@ mod tests { create_on_compaction: Mode::Disable, ..Default::default() }, - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1851,8 +1764,6 @@ mod tests { with_inverted: false, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1864,8 +1775,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1878,8 +1787,6 @@ mod tests { with_inverted: true, with_fulltext: false, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1891,8 +1798,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1905,8 +1810,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: false, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1918,8 +1821,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1939,8 +1840,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1952,8 +1851,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(0)) .await; @@ -1961,81 +1858,6 @@ mod tests { assert!(indexer.inverted_indexer.is_some()); } - #[cfg(feature = "vector_index")] - #[tokio::test] - async fn test_update_flat_builds_vector_index() { - use datatypes::arrow::array::BinaryBuilder; - use datatypes::arrow::datatypes::{DataType, Field, Schema}; - - struct TestPathProvider; - - impl FilePathProvider for TestPathProvider { - fn build_index_file_path(&self, file_id: RegionFileId) -> String { - format!("index/{}.puffin", file_id) - } - - fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String { - format!("index/{}.puffin", index_id) - } - - fn build_sst_file_path(&self, file_id: RegionFileId) -> String { - format!("sst/{}.parquet", file_id) - } - } - - fn f32s_to_bytes(values: &[f32]) -> Vec { - let mut bytes = Vec::with_capacity(values.len() * 4); - for v in values { - bytes.extend_from_slice(&v.to_le_bytes()); - } - bytes - } - - let (dir, factory) = - PuffinManagerFactory::new_for_test_async("test_update_flat_builds_vector_index_").await; - let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await; - - let metadata = mock_region_metadata(MetaConfig { - with_inverted: false, - with_fulltext: false, - with_skipping_bloom: false, - with_vector: true, - }); - - let mut indexer = IndexerBuilderImpl { - build_type: IndexBuildType::Flush, - metadata, - puffin_manager: factory.build(mock_object_store(), TestPathProvider), - write_cache_enabled: false, - intermediate_manager: intm_manager, - index_options: IndexOptions::default(), - inverted_index_config: InvertedIndexConfig::default(), - fulltext_index_config: FulltextIndexConfig::default(), - bloom_filter_index_config: BloomFilterConfig::default(), - vector_index_config: Default::default(), - } - .build(random_region_file_id(), 0, Some(1024)) - .await; - - assert!(indexer.vector_indexer.is_some()); - - let vec1 = f32s_to_bytes(&[1.0, 0.0, 0.0, 0.0]); - let vec2 = f32s_to_bytes(&[0.0, 1.0, 0.0, 0.0]); - - let mut builder = BinaryBuilder::with_capacity(2, vec1.len() + vec2.len()); - builder.append_value(&vec1); - builder.append_value(&vec2); - - let schema = Arc::new(Schema::new(vec![Field::new("vec", DataType::Binary, true)])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(builder.finish())]).unwrap(); - - indexer.update_flat(&batch).await; - let output = indexer.finish().await; - - assert!(output.vector_index.is_available()); - assert!(output.vector_index.columns.contains(&6)); - } - #[tokio::test] async fn test_index_build_task_sst_not_exist() { let env = SchedulerEnv::new().await; @@ -2446,8 +2268,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }); let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await; diff --git a/src/mito2/src/sst/index/bloom_filter/applier.rs b/src/mito2/src/sst/index/bloom_filter/applier.rs index e2d25e4bf82..1faae466c89 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier.rs @@ -422,23 +422,14 @@ impl BloomFilterIndexApplier { reader: R, predicates: &[InListPredicate], output: &mut [(usize, Vec>)], - mut metrics: Option<&mut BloomFilterIndexApplyMetrics>, + metrics: Option<&mut BloomFilterIndexApplyMetrics>, ) -> std::result::Result<(), index::bloom_filter::error::Error> { let mut applier = BloomFilterApplier::new(Box::new(reader)).await?; - - for (_, row_group_output) in output.iter_mut() { - // All rows are filtered out, skip the search - if row_group_output.is_empty() { - continue; - } - - let read_metrics = metrics.as_deref_mut().map(|m| &mut m.read_metrics); - *row_group_output = applier - .search(predicates, row_group_output, read_metrics) - .await?; - } - - Ok(()) + let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::>(); + let read_metrics = metrics.map(|m| &mut m.read_metrics); + applier + .search_groups(predicates, &mut row_groups, read_metrics) + .await } /// Returns compatible bloom filter predicates with the given SST metadata. diff --git a/src/mito2/src/sst/index/column_test.rs b/src/mito2/src/sst/index/column_test.rs index 146bed120c0..fb7677eff58 100644 --- a/src/mito2/src/sst/index/column_test.rs +++ b/src/mito2/src/sst/index/column_test.rs @@ -29,7 +29,9 @@ use index::bitmap::{Bitmap, BitmapType}; use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl}; use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader}; use mito_codec::index::IndexValueCodec; -use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; +use mito_codec::row_converter::{ + DensePrimaryKeyCodec, PrimaryKeyCodec, PrimaryKeyCodecExt, SortField, +}; use object_store::ObjectStore; use object_store::services::Memory; use prost::Message; @@ -197,17 +199,87 @@ async fn materialized_runs_match_legacy_index_bytes_and_bloom_lookups() { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut expected = None; - // The legacy Batch path is an independent whole-PK decode/owned-encoding oracle. + // Demand is the union of active creators' PK columns, not their counts or + // all columns in the metadata. A malformed unneeded suffix distinguishes + // the lazy path from a premature whole-key decode. + for (inverted, bloom, disable_bloom, all_pk) in [ + (vec![0, 1], vec![], false, true), + (vec![], vec![0, 1], false, true), + (vec![0], vec![1], false, true), + (vec![0, 2], vec![0, 2], false, false), + (vec![0], vec![1], true, false), + (vec![2], vec![2], false, false), + ] { + let mut case_metadata = RegionMetadataBuilder::new(metadata.region_id); + for column in &metadata.column_metadatas { + let mut column = column.clone(); + column + .column_schema + .set_inverted_index(inverted.contains(&column.column_id)); + if !bloom.contains(&column.column_id) { + column.column_schema.unset_skipping_options().unwrap(); + } + case_metadata.push_column_metadata(column); + } + case_metadata.primary_key(metadata.primary_key.clone()); + let mut case_builder = builder.clone(); + case_builder.metadata = Arc::new(case_metadata.build().unwrap()); + if disable_bloom { + case_builder.bloom_filter_index_config.create_on_flush = crate::config::Mode::Disable; + } + let mut indexer = case_builder + .build( + RegionFileId::new(metadata.region_id, FileId::random()), + 0, + None, + ) + .await; + let make_batch = |key| { + let mut batch_builder = BatchBuilder::new(key); + batch_builder + .timestamps_array(batch.column(3).slice(0, 1)) + .unwrap(); + batch_builder + .sequences_array(batch.column(5).slice(0, 1)) + .unwrap(); + batch_builder + .op_types_array(batch.column(6).slice(0, 1)) + .unwrap(); + batch_builder.build().unwrap() + }; + let mut valid = make_batch(encoded[0].clone()); + indexer.prepare_primary_key(&mut valid).unwrap(); + if all_pk { + assert_eq!( + valid.pk_values(), + Some( + &DensePrimaryKeyCodec::new(&metadata) + .decode(&encoded[0]) + .unwrap() + ) + ); + } else { + assert!(valid.pk_values().is_none()); + } + let mut truncated = make_batch(vec![0, 1]); + assert_eq!(indexer.prepare_primary_key(&mut truncated).is_err(), all_pk); + assert!(truncated.pk_values().is_none()); + indexer.abort().await; + // Aborted creators must not retain their previous all-PK demand. + indexer.prepare_primary_key(&mut truncated).unwrap(); + } + // Explicit full decoding is the oracle for lazy encoded-only Batch inputs. // One-row slices disable run merging; a tag-only projection also exercises the no-PK fallback. - for mode in ["legacy", "runs", "single_rows", "no_pk"] { + for mode in ["eager", "legacy", "lazy", "runs", "single_rows", "no_pk"] { let file = RegionFileId::new(metadata.region_id, FileId::random()); let mut indexer = builder.build(file, 0, None).await; + if mode == "lazy" { + indexer.dense_pk_decoder = None; + } match mode { - "legacy" => { + "eager" | "legacy" | "lazy" => { for (row, key) in encoded.iter().enumerate() { let mut old = BatchBuilder::new(key.clone()); old.push_field_array(2, batch.column(2).slice(row, 1)) @@ -215,7 +287,13 @@ async fn materialized_runs_match_legacy_index_bytes_and_bloom_lookups() { old.timestamps_array(batch.column(3).slice(row, 1)).unwrap(); old.sequences_array(batch.column(5).slice(row, 1)).unwrap(); old.op_types_array(batch.column(6).slice(row, 1)).unwrap(); - indexer.update(&mut old.build().unwrap()).await; + let mut old = old.build().unwrap(); + if mode == "eager" { + old.set_pk_values( + DensePrimaryKeyCodec::new(&metadata).decode(key).unwrap(), + ); + } + indexer.update(&mut old).await; } } "single_rows" => { diff --git a/src/mito2/src/sst/index/fulltext_index/applier.rs b/src/mito2/src/sst/index/fulltext_index/applier.rs index 5ccf4e57a4b..4efb2a5ccb9 100644 --- a/src/mito2/src/sst/index/fulltext_index/applier.rs +++ b/src/mito2/src/sst/index/fulltext_index/applier.rs @@ -464,23 +464,15 @@ impl FulltextIndexApplier { let mut applier = BloomFilterApplier::new(reader) .await .context(ApplyBloomFilterIndexSnafu)?; - for (_, row_group_output) in output.iter_mut() { - // All rows are filtered out, skip the search - if row_group_output.is_empty() { - continue; - } - - *row_group_output = applier - .search( - &predicates, - row_group_output, - metrics - .as_deref_mut() - .map(|m| &mut m.bloom_filter_read_metrics), - ) - .await - .context(ApplyBloomFilterIndexSnafu)?; - } + let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::>(); + applier + .search_groups( + &predicates, + &mut row_groups, + metrics.map(|m| &mut m.bloom_filter_read_metrics), + ) + .await + .context(ApplyBloomFilterIndexSnafu)?; Ok(true) } diff --git a/src/mito2/src/sst/index/indexer/abort.rs b/src/mito2/src/sst/index/indexer/abort.rs index a928eb186da..36c002f13f0 100644 --- a/src/mito2/src/sst/index/indexer/abort.rs +++ b/src/mito2/src/sst/index/indexer/abort.rs @@ -20,11 +20,10 @@ use crate::sst::index::Indexer; impl Indexer { pub(crate) async fn do_abort(&mut self) { + self.dense_pk_decoder = None; self.do_abort_inverted_index().await; self.do_abort_fulltext_index().await; self.do_abort_bloom_filter().await; - #[cfg(feature = "vector_index")] - self.do_abort_vector_index().await; self.do_prune_intm_sst_dir().await; if self.write_cache_enabled { self.do_abort_clean_fs_temp_dir().await; @@ -108,26 +107,4 @@ impl Indexer { .to_string(); TempFileCleaner::clean_atomic_dir_files(fs_accessor.store().store(), &[&fs_handle]).await; } - - #[cfg(feature = "vector_index")] - async fn do_abort_vector_index(&mut self) { - let Some(mut indexer) = self.vector_indexer.take() else { - return; - }; - let Err(err) = indexer.abort().await else { - return; - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to abort vector index, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to abort vector index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - } } diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index c7966823ddf..6a6c3f34d0b 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -17,8 +17,6 @@ use puffin::puffin_manager::{PuffinManager, PuffinWriter}; use store_api::storage::ColumnId; use crate::sst::file::{RegionFileId, RegionIndexId}; -#[cfg(feature = "vector_index")] -use crate::sst::index::VectorIndexOutput; use crate::sst::index::puffin_manager::SstPuffinWriter; use crate::sst::index::statistics::{ByteCount, RowCount}; use crate::sst::index::{ @@ -27,6 +25,7 @@ use crate::sst::index::{ impl Indexer { pub(crate) async fn do_finish(&mut self) -> IndexOutput { + self.dense_pk_decoder = None; let mut output = IndexOutput::default(); let Some(mut writer) = self.build_puffin_writer().await else { @@ -56,15 +55,6 @@ impl Indexer { return IndexOutput::default(); } - #[cfg(feature = "vector_index")] - { - let success = self.do_finish_vector_index(&mut writer, &mut output).await; - if !success { - self.do_abort().await; - return IndexOutput::default(); - } - } - self.do_prune_intm_sst_dir().await; output.file_size = self.do_finish_puffin_writer(writer).await; output.version = self.index_version; @@ -287,63 +277,6 @@ impl Indexer { output.columns = column_ids; } - #[cfg(feature = "vector_index")] - async fn do_finish_vector_index( - &mut self, - puffin_writer: &mut SstPuffinWriter, - index_output: &mut IndexOutput, - ) -> bool { - let Some(mut indexer) = self.vector_indexer.take() else { - return true; - }; - - let column_ids = indexer.column_ids().collect(); - let err = match indexer.finish(puffin_writer).await { - Ok((row_count, byte_count)) => { - self.fill_vector_index_output( - &mut index_output.vector_index, - row_count, - byte_count, - column_ids, - ); - return true; - } - Err(err) => err, - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to finish vector index, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to finish vector index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - - false - } - - #[cfg(feature = "vector_index")] - fn fill_vector_index_output( - &mut self, - output: &mut VectorIndexOutput, - row_count: RowCount, - byte_count: ByteCount, - column_ids: Vec, - ) { - debug!( - "Vector index created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}", - self.region_id, self.file_id, byte_count, row_count, column_ids - ); - - output.index_size = byte_count; - output.row_count = row_count; - output.columns = column_ids; - } - pub(crate) async fn do_prune_intm_sst_dir(&mut self) { if let Some(manager) = self.intermediate_manager.take() && let Err(e) = manager.prune_sst_dir(&self.region_id, &self.file_id).await diff --git a/src/mito2/src/sst/index/indexer/update.rs b/src/mito2/src/sst/index/indexer/update.rs index 095da337f49..7e2e1e35022 100644 --- a/src/mito2/src/sst/index/indexer/update.rs +++ b/src/mito2/src/sst/index/indexer/update.rs @@ -24,6 +24,11 @@ impl Indexer { return; } + if !self.do_prepare_primary_key(batch) { + self.do_abort().await; + return; + } + if !self.do_update_inverted_index(batch).await { self.do_abort().await; } @@ -33,10 +38,23 @@ impl Indexer { if !self.do_update_bloom_filter(batch).await { self.do_abort().await; } - #[cfg(feature = "vector_index")] - if !self.do_update_vector_index(batch).await { - self.do_abort().await; + } + + /// Handles decode errors before entering asynchronous cleanup, following the + /// creators' update policy without carrying a decode error across an await. + fn do_prepare_primary_key(&self, batch: &mut Batch) -> bool { + let Err(err) = self.prepare_primary_key(batch) else { + return true; + }; + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to decode primary key for indexes, region_id: {}, file_id: {}, err: {:?}", + self.region_id, self.file_id, err + ); + } else { + warn!(err; "Failed to decode primary key for indexes, region_id: {}, file_id: {}", self.region_id, self.file_id); } + false } /// Returns false if the update failed. @@ -114,32 +132,6 @@ impl Indexer { false } - /// Returns false if the update failed. - #[cfg(feature = "vector_index")] - async fn do_update_vector_index(&mut self, batch: &mut Batch) -> bool { - let Some(creator) = self.vector_indexer.as_mut() else { - return true; - }; - - let Err(err) = creator.update(batch).await else { - return true; - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to update vector index, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to update vector index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - - false - } - pub(crate) async fn do_update_flat(&mut self, batch: &RecordBatch) { if batch.num_rows() == 0 { return; @@ -154,10 +146,6 @@ impl Indexer { if !self.do_update_flat_bloom_filter(batch).await { self.do_abort().await; } - #[cfg(feature = "vector_index")] - if !self.do_update_flat_vector_index(batch).await { - self.do_abort().await; - } } /// Returns false if the update failed. @@ -234,30 +222,4 @@ impl Indexer { false } - - /// Returns false if the update failed. - #[cfg(feature = "vector_index")] - async fn do_update_flat_vector_index(&mut self, batch: &RecordBatch) -> bool { - let Some(creator) = self.vector_indexer.as_mut() else { - return true; - }; - - let Err(err) = creator.update_flat(batch).await else { - return true; - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to update vector index with flat format, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to update vector index with flat format, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - - false - } } diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 9ea5258ee3b..d1dee226515 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -204,10 +204,12 @@ impl InvertedIndexer { &mut self.value_buf, ) .context(EncodeSnafu)?; - self.index_creator - .push_with_name_n(target_key, elem, count) - .await - .context(PushIndexValueSnafu)?; + if self.index_creator.push_with_name_n(target_key, elem, count) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } } else if is_sparse && column_meta.semantic_type == SemanticType::Tag { if self.codec.pk_col_info(*col_id).is_some() { @@ -233,10 +235,15 @@ impl InvertedIndexer { &mut self.value_buf, ) .context(DecodeSnafu)?; - self.index_creator + if self + .index_creator .push_with_name_n(target_key, value, count) - .await - .context(PushIndexValueSnafu)?; + { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } } } @@ -307,10 +314,12 @@ impl InvertedIndexer { }) .transpose()?; - self.index_creator - .push_with_name_n(target_key, value, n) - .await - .context(PushIndexValueSnafu)?; + if self.index_creator.push_with_name_n(target_key, value, n) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } // fields None => { @@ -326,10 +335,12 @@ impl InvertedIndexer { self.value_buf.clear(); let value = values.data.get_ref(i); if value.is_null() { - self.index_creator - .push_with_name(target_key, None) - .await - .context(PushIndexValueSnafu)?; + if self.index_creator.push_with_name(target_key, None) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } else { IndexValueCodec::encode_nonnull_value( value, @@ -337,10 +348,15 @@ impl InvertedIndexer { &mut self.value_buf, ) .context(EncodeSnafu)?; - self.index_creator + if self + .index_creator .push_with_name(target_key, Some(&self.value_buf)) - .await - .context(PushIndexValueSnafu)?; + { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } } } diff --git a/src/mito2/src/sst/index/sparse_test.rs b/src/mito2/src/sst/index/sparse_test.rs index 15d26bd1bc7..c6031d50ad6 100644 --- a/src/mito2/src/sst/index/sparse_test.rs +++ b/src/mito2/src/sst/index/sparse_test.rs @@ -141,8 +141,6 @@ async fn sparse_and_materialized_tags_produce_identical_indexes() { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut expected = None; for batch in [materialized, sparse, mixed] { diff --git a/src/mito2/src/sst/index/store.rs b/src/mito2/src/sst/index/store.rs index e737252652d..0a99744986a 100644 --- a/src/mito2/src/sst/index/store.rs +++ b/src/mito2/src/sst/index/store.rs @@ -262,6 +262,9 @@ impl AsyncWrite for InstrumentedAsyncWrite<'_, W> } } +/// Ranges fetched in parallel by one `read_vec`, as for parquet reads. +const FETCH_CONCURRENCY: usize = 8; + /// Implements `RangeReader` for `ObjectStore` and record metrics. pub(crate) struct InstrumentedRangeReader<'a> { store: ObjectStore, @@ -309,9 +312,16 @@ impl RangeReader for InstrumentedRangeReader<'_> { } async fn read_vec(&self, ranges: &[Range]) -> io::Result> { + // OpenDAL merges ranges less than 1 MiB apart by default without bounding the merged + // length, so sparse page reads could fetch most of the file. Only merge adjacent + // ranges, and fetch the rest in parallel like parquet reads do. + // TODO: decide whether to merge across small gaps based on object store latency, + // bytes read and query latency. let bufs = self .store - .reader(&self.path) + .reader_with(&self.path) + .gap(0) + .concurrent(FETCH_CONCURRENCY) .await? .fetch(ranges.to_owned()) .await?; diff --git a/src/mito2/src/sst/index/vector_index.rs b/src/mito2/src/sst/index/vector_index.rs deleted file mode 100644 index f1b9aa3899e..00000000000 --- a/src/mito2/src/sst/index/vector_index.rs +++ /dev/null @@ -1,24 +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. - -//! Vector index module for HNSW-based approximate nearest neighbor search. -//! -//! This module provides the mito2 integration layer for vector indexes. -//! The core algorithms (creator, reader, applier) are in the `index` crate. - -pub(crate) mod applier; -pub(crate) mod creator; - -/// The blob type identifier for vector index in puffin files. -pub(crate) const INDEX_BLOB_TYPE: &str = "greptime-vector-index-v1"; diff --git a/src/mito2/src/sst/index/vector_index/applier.rs b/src/mito2/src/sst/index/vector_index/applier.rs deleted file mode 100644 index 40c8c257277..00000000000 --- a/src/mito2/src/sst/index/vector_index/applier.rs +++ /dev/null @@ -1,503 +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. - -//! Vector index applier for KNN search. - -use std::sync::Arc; - -use common_base::range_read::RangeReader; -use common_telemetry::warn; -use index::vector::VectorDistanceMetric; -use index::vector::apply::HnswVectorIndexApplier; -use puffin::puffin_manager::cache::PuffinMetadataCacheRef; -use puffin::puffin_manager::{PuffinManager, PuffinReader}; -use snafu::ResultExt; -use store_api::storage::ColumnId; - -use crate::access_layer::{RegionFilePathFactory, WriteCachePathProvider}; -use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; -use crate::cache::index::vector_index::{ - CachedVectorIndex, VectorIndexCacheKey, VectorIndexCacheRef, -}; -use crate::error::{ApplyVectorIndexSnafu, PuffinBuildReaderSnafu, PuffinReadBlobSnafu, Result}; -use crate::sst::file::RegionIndexId; -use crate::sst::index::puffin_manager::{BlobReader, PuffinManagerFactory}; -use crate::sst::index::trigger_index_background_download; -use crate::sst::index::vector_index::INDEX_BLOB_TYPE; - -/// Result of applying vector index. -#[derive(Debug)] -pub struct VectorIndexApplyOutput { - /// Row offsets in the SST file. - pub row_offsets: Vec, - /// Distances to the query vector. - /// TODO(dennis): Not used yet, reserved for future use. - #[allow(dead_code)] - pub distances: Vec, -} - -/// Vector index applier for KNN search against SST blobs. -pub struct VectorIndexApplier { - table_dir: String, - path_type: store_api::region_request::PathType, - object_store: object_store::ObjectStore, - puffin_manager_factory: PuffinManagerFactory, - file_cache: Option, - puffin_metadata_cache: Option, - vector_index_cache: Option, - column_id: ColumnId, - query_vector: Vec, - metric: VectorDistanceMetric, -} - -pub type VectorIndexApplierRef = Arc; - -impl VectorIndexApplier { - pub fn new( - table_dir: String, - path_type: store_api::region_request::PathType, - object_store: object_store::ObjectStore, - puffin_manager_factory: PuffinManagerFactory, - column_id: ColumnId, - query_vector: Vec, - metric: VectorDistanceMetric, - ) -> Self { - Self { - table_dir, - path_type, - object_store, - puffin_manager_factory, - file_cache: None, - puffin_metadata_cache: None, - vector_index_cache: None, - column_id, - query_vector, - metric, - } - } - - pub fn with_file_cache(mut self, file_cache: Option) -> Self { - self.file_cache = file_cache; - self - } - - pub fn with_puffin_metadata_cache( - mut self, - puffin_metadata_cache: Option, - ) -> Self { - self.puffin_metadata_cache = puffin_metadata_cache; - self - } - - pub fn with_vector_index_cache(mut self, cache: Option) -> Self { - self.vector_index_cache = cache; - self - } - - /// Applies vector index to the file and returns candidates. - /// - /// This method loads the vector index blob (from cache or remote), runs - /// a KNN search against the indexed vectors, and maps the HNSW keys back - /// to row offsets in the SST file. It returns only row offsets; callers - /// are responsible for any higher-level ordering or limit enforcement. - pub async fn apply_with_k( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - k: usize, - ) -> Result { - if k == 0 { - return Ok(VectorIndexApplyOutput { - row_offsets: Vec::new(), - distances: Vec::new(), - }); - } - - let cached = self.load_or_read_index(file_id, file_size_hint).await?; - let Some(cached) = cached else { - return Ok(VectorIndexApplyOutput { - row_offsets: Vec::new(), - distances: Vec::new(), - }); - }; - - let applier = cached.applier(); - if self.query_vector.len() != applier.dimensions() as usize { - return ApplyVectorIndexSnafu { - reason: format!( - "Query vector dimension {} does not match index dimension {}", - self.query_vector.len(), - applier.dimensions() - ), - } - .fail(); - } - if self.metric != applier.metric() { - return ApplyVectorIndexSnafu { - reason: format!( - "Query metric {} does not match index metric {}", - self.metric, - applier.metric() - ), - } - .fail(); - } - - let output = applier.search(&self.query_vector, k).map_err(|e| { - ApplyVectorIndexSnafu { - reason: e.to_string(), - } - .build() - })?; - - Ok(VectorIndexApplyOutput { - row_offsets: output.row_offsets, - distances: output.distances, - }) - } - - async fn load_or_read_index( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - ) -> Result>> { - let cache_key = - VectorIndexCacheKey::new(file_id.file_id(), file_id.version, self.column_id); - if let Some(cache) = &self.vector_index_cache - && let Some(cached) = cache.get(&cache_key) - { - return Ok(Some(cached)); - } - - let reader = match self.cached_blob_reader(file_id, file_size_hint).await { - Ok(Some(reader)) => reader, - Ok(None) => self.remote_blob_reader(file_id, file_size_hint).await?, - Err(err) => { - if is_blob_not_found(&err) { - self.remote_blob_reader(file_id, file_size_hint).await? - } else { - warn!(err; "Failed to read cached vector index blob, fallback to remote"); - self.remote_blob_reader(file_id, file_size_hint).await? - } - } - }; - - let blob_data = read_all_blob(reader, file_size_hint).await?; - if blob_data.is_empty() { - return Ok(None); - } - - let applier = HnswVectorIndexApplier::from_blob(&blob_data).map_err(|e| { - ApplyVectorIndexSnafu { - reason: e.to_string(), - } - .build() - })?; - - let cached = Arc::new(CachedVectorIndex::new(Box::new(applier))); - if let Some(cache) = &self.vector_index_cache { - cache.insert(cache_key, cached.clone()); - } - - Ok(Some(cached)) - } - - async fn cached_blob_reader( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - ) -> Result> { - let Some(file_cache) = &self.file_cache else { - return Ok(None); - }; - - let index_key = IndexKey::new( - file_id.region_id(), - file_id.file_id(), - FileType::Puffin(file_id.version), - ); - if file_cache.get(index_key).await.is_none() { - return Ok(None); - } - - let puffin_manager = self.puffin_manager_factory.build( - file_cache.local_store(), - WriteCachePathProvider::new(file_cache.clone()), - ); - let blob_name = column_blob_name(self.column_id); - - let reader = puffin_manager - .reader(&file_id) - .await - .context(PuffinBuildReaderSnafu)? - .with_file_size_hint(file_size_hint) - .blob(&blob_name) - .await - .context(PuffinReadBlobSnafu)? - .reader() - .await - .context(PuffinBuildReaderSnafu)?; - Ok(Some(reader)) - } - - async fn remote_blob_reader( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - ) -> Result { - let path_factory = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type); - - trigger_index_background_download( - self.file_cache.as_ref(), - &file_id, - file_size_hint, - &path_factory, - &self.object_store, - ); - - let puffin_manager = self - .puffin_manager_factory - .build(self.object_store.clone(), path_factory) - .with_puffin_metadata_cache(self.puffin_metadata_cache.clone()); - - let blob_name = column_blob_name(self.column_id); - - puffin_manager - .reader(&file_id) - .await - .context(PuffinBuildReaderSnafu)? - .with_file_size_hint(file_size_hint) - .blob(&blob_name) - .await - .context(PuffinReadBlobSnafu)? - .reader() - .await - .context(PuffinBuildReaderSnafu) - } -} - -fn column_blob_name(column_id: ColumnId) -> String { - format!("{INDEX_BLOB_TYPE}-{}", column_id) -} - -fn is_blob_not_found(err: &crate::error::Error) -> bool { - matches!( - err, - crate::error::Error::PuffinReadBlob { - source: puffin::error::Error::BlobNotFound { .. }, - .. - } - ) -} - -async fn read_all_blob(reader: BlobReader, file_size_hint: Option) -> Result> { - let metadata = reader.metadata().await.map_err(|e| { - ApplyVectorIndexSnafu { - reason: format!("Failed to read vector index metadata: {}", e), - } - .build() - })?; - if let Some(limit) = file_size_hint - && metadata.content_length > limit - { - return ApplyVectorIndexSnafu { - reason: format!( - "Vector index blob size {} exceeds file size hint {}", - metadata.content_length, limit - ), - } - .fail(); - } - let bytes = reader.read(0..metadata.content_length).await.map_err(|e| { - ApplyVectorIndexSnafu { - reason: format!("Failed to read vector index data: {}", e), - } - .build() - })?; - Ok(bytes.to_vec()) -} - -#[cfg(test)] -mod tests { - use common_test_util::temp_dir::TempDir; - use futures::io::Cursor; - use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; - use index::vector::engine::{self, VectorIndexConfig}; - use index::vector::format::{distance_metric_to_proto, engine_type_to_proto}; - use object_store::ObjectStore; - use object_store::services::Memory; - use prost::Message; - use puffin::puffin_manager::PuffinWriter; - use roaring::RoaringBitmap; - use store_api::region_request::PathType; - use store_api::storage::{ColumnId, FileId, VectorIndexEngineType}; - - use super::*; - use crate::access_layer::RegionFilePathFactory; - use crate::sst::file::RegionFileId; - use crate::sst::index::puffin_manager::PuffinManagerFactory; - - async fn build_applier_with_blob( - blob: Vec, - column_id: ColumnId, - query_vector: Vec, - metric: VectorDistanceMetric, - ) -> (TempDir, VectorIndexApplier, RegionIndexId, u64) { - let (dir, puffin_manager_factory) = - PuffinManagerFactory::new_for_test_async("test_vector_index_applier_").await; - let object_store = ObjectStore::new(Memory::default()).unwrap(); - let file_id = RegionFileId::new(0.into(), FileId::random()); - let index_id = RegionIndexId::new(file_id, 0); - let table_dir = "table_dir".to_string(); - - let puffin_manager = puffin_manager_factory.build( - object_store.clone(), - RegionFilePathFactory::new(table_dir.clone(), PathType::Bare), - ); - let mut writer = puffin_manager.writer(&index_id).await.unwrap(); - let blob_name = column_blob_name(column_id); - let _bytes_written = writer - .put_blob( - blob_name.as_str(), - Cursor::new(blob), - Default::default(), - Default::default(), - ) - .await - .unwrap(); - let file_size = writer.finish().await.unwrap(); - - let applier = VectorIndexApplier::new( - table_dir, - PathType::Bare, - object_store, - puffin_manager_factory, - column_id, - query_vector, - metric, - ); - - (dir, applier, index_id, file_size) - } - - fn build_blob_with_vectors( - config: &VectorIndexConfig, - vectors: Vec<(u64, Vec)>, - null_bitmap: &RoaringBitmap, - total_rows: u64, - indexed_rows: u64, - ) -> Vec { - let mut engine = engine::create_engine(config.engine, config).unwrap(); - for (key, vector) in vectors { - engine.add(key, &vector).unwrap(); - } - let index_size = engine.serialized_length(); - let mut index_bytes = vec![0u8; index_size]; - engine.save_to_buffer(&mut index_bytes).unwrap(); - - let mut null_bitmap_bytes = Vec::new(); - null_bitmap.serialize_into(&mut null_bitmap_bytes).unwrap(); - - let null_count = total_rows - indexed_rows; - let meta = VectorIndexMeta { - engine: engine_type_to_proto(config.engine).into(), - dim: config.dim as u32, - metric: distance_metric_to_proto(config.distance_metric).into(), - connectivity: config.connectivity as u32, - expansion_add: config.expansion_add as u32, - expansion_search: config.expansion_search as u32, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_bytes.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: total_rows, - indexed_row_count: indexed_rows, - null_count, - }), - }; - let meta_bytes = meta.encode_to_vec(); - let meta_size = meta_bytes.len() as u32; - - // Build blob: null_bitmap | index_data | meta | meta_size - let mut blob = Vec::new(); - blob.extend_from_slice(&null_bitmap_bytes); - blob.extend_from_slice(&index_bytes); - blob.extend_from_slice(&meta_bytes); - blob.extend_from_slice(&meta_size.to_le_bytes()); - blob - } - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 2, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[tokio::test] - async fn test_apply_with_k_returns_offsets() { - let config = test_config(); - let mut null_bitmap = RoaringBitmap::new(); - null_bitmap.insert(1); - let blob = build_blob_with_vectors( - &config, - vec![(0, vec![1.0, 0.0]), (1, vec![0.0, 1.0])], - &null_bitmap, - 3, - 2, - ); - let (_dir, applier, index_id, size_bytes) = - build_applier_with_blob(blob, 1, vec![1.0, 0.0], VectorDistanceMetric::L2sq).await; - let output = applier - .apply_with_k(index_id, Some(size_bytes), 2) - .await - .unwrap(); - assert_eq!(output.row_offsets, vec![0, 2]); - } - - #[tokio::test] - async fn test_apply_with_k_dimension_mismatch() { - let config = test_config(); - let null_bitmap = RoaringBitmap::new(); - let blob = build_blob_with_vectors(&config, vec![(0, vec![1.0, 0.0])], &null_bitmap, 1, 1); - let (_dir, applier, index_id, size_bytes) = - build_applier_with_blob(blob, 1, vec![1.0, 0.0, 0.0], VectorDistanceMetric::L2sq).await; - let res = applier.apply_with_k(index_id, Some(size_bytes), 1).await; - assert!(res.is_err()); - } - - #[tokio::test] - async fn test_apply_with_k_empty_blob() { - let config = VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 1, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - }; - let null_bitmap = RoaringBitmap::new(); - let blob = build_blob_with_vectors(&config, Vec::new(), &null_bitmap, 0, 0); - let (_dir, applier, index_id, size_bytes) = - build_applier_with_blob(blob, 1, vec![1.0], VectorDistanceMetric::L2sq).await; - let output = applier - .apply_with_k(index_id, Some(size_bytes), 1) - .await - .unwrap(); - assert!(output.row_offsets.is_empty()); - } -} diff --git a/src/mito2/src/sst/index/vector_index/creator.rs b/src/mito2/src/sst/index/vector_index/creator.rs deleted file mode 100644 index 2971c12cc66..00000000000 --- a/src/mito2/src/sst/index/vector_index/creator.rs +++ /dev/null @@ -1,536 +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. - -//! Vector indexer for managing vector indexes across multiple columns. - -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; - -use common_telemetry::warn; -use datatypes::arrow::array::{Array, BinaryArray}; -use datatypes::arrow::record_batch::RecordBatch; -use datatypes::data_type::ConcreteDataType; -use datatypes::prelude::ValueRef; -use index::vector::VectorIndexOptions; -use index::vector::create::{HnswVectorIndexCreator, IndexCreator}; -use index::vector::engine::VectorIndexConfig; -use index::vector::util::bytes_to_f32_slice; -use puffin::puffin_manager::PutOptions; -use snafu::ensure; -use store_api::metadata::RegionMetadataRef; -use store_api::storage::{ColumnId, FileId}; - -use crate::error::{ - OperateAbortedIndexSnafu, Result, VectorIndexBuildSnafu, VectorIndexFinishSnafu, -}; -use crate::metrics::{INDEX_CREATE_BYTES_TOTAL, INDEX_CREATE_ROWS_TOTAL}; -use crate::read::Batch; -use crate::sst::index::TYPE_VECTOR_INDEX; -use crate::sst::index::intermediate::{ - IntermediateLocation, IntermediateManager, TempFileProvider, -}; -use crate::sst::index::puffin_manager::SstPuffinWriter; -use crate::sst::index::statistics::{ByteCount, RowCount, Statistics}; -use crate::sst::index::vector_index::INDEX_BLOB_TYPE; - -/// Helper to create VectorIndexConfig from VectorIndexOptions. -fn create_config(dim: usize, options: &VectorIndexOptions) -> VectorIndexConfig { - VectorIndexConfig { - engine: options.engine, - dim, - distance_metric: options.metric, - connectivity: options.connectivity as usize, - expansion_add: options.expansion_add as usize, - expansion_search: options.expansion_search as usize, - } -} - -fn build_err(e: impl std::fmt::Display) -> crate::error::Error { - VectorIndexBuildSnafu { - reason: e.to_string(), - } - .build() -} - -/// The indexer for vector indexes across multiple columns. -pub struct VectorIndexer { - /// Per-column vector index creators. - creators: HashMap, - /// Provider for intermediate files. - temp_file_provider: Arc, - /// Whether the indexing process has been aborted. - aborted: bool, - /// Statistics for this indexer. - stats: Statistics, - /// Global memory usage tracker. - #[allow(dead_code)] - global_memory_usage: Arc, - /// Region metadata for column lookups. - metadata: RegionMetadataRef, - /// Memory usage threshold. - memory_usage_threshold: Option, -} - -impl VectorIndexer { - /// Creates a new vector indexer. - /// - /// Returns `None` if there are no vector columns that need indexing. - pub fn new( - sst_file_id: FileId, - metadata: &RegionMetadataRef, - intermediate_manager: IntermediateManager, - memory_usage_threshold: Option, - vector_index_options: &HashMap, - ) -> Result> { - let mut creators = HashMap::new(); - - let temp_file_provider = Arc::new(TempFileProvider::new( - IntermediateLocation::new(&metadata.region_id, &sst_file_id), - intermediate_manager, - )); - let global_memory_usage = Arc::new(AtomicUsize::new(0)); - - // Find all vector columns that have vector index enabled - for column in &metadata.column_metadatas { - // Check if this column has vector index options configured - let Some(options) = vector_index_options.get(&column.column_id) else { - continue; - }; - - // Verify the column is a vector type - let ConcreteDataType::Vector(vector_type) = &column.column_schema.data_type else { - continue; - }; - - let config = create_config(vector_type.dim as usize, options); - let creator = HnswVectorIndexCreator::new(config).map_err(build_err)?; - creators.insert(column.column_id, creator); - } - - if creators.is_empty() { - return Ok(None); - } - - let indexer = Self { - creators, - temp_file_provider, - aborted: false, - stats: Statistics::new(TYPE_VECTOR_INDEX), - global_memory_usage, - metadata: metadata.clone(), - memory_usage_threshold, - }; - - Ok(Some(indexer)) - } - - /// Updates index with a batch of rows. - /// Garbage will be cleaned up if failed to update. - pub async fn update(&mut self, batch: &mut Batch) -> Result<()> { - ensure!(!self.aborted, OperateAbortedIndexSnafu); - - if self.creators.is_empty() { - return Ok(()); - } - - if let Err(update_err) = self.do_update(batch).await { - // Clean up garbage if failed to update - if let Err(err) = self.do_cleanup().await { - if cfg!(any(test, feature = "test")) { - panic!("Failed to clean up vector index creator, err: {err:?}"); - } else { - warn!(err; "Failed to clean up vector index creator"); - } - } - return Err(update_err); - } - - Ok(()) - } - - /// Updates index with a flat format `RecordBatch`. - /// Garbage will be cleaned up if failed to update. - pub async fn update_flat(&mut self, batch: &RecordBatch) -> Result<()> { - ensure!(!self.aborted, OperateAbortedIndexSnafu); - - if self.creators.is_empty() || batch.num_rows() == 0 { - return Ok(()); - } - - if let Err(update_err) = self.do_update_flat(batch).await { - // Clean up garbage if failed to update - if let Err(err) = self.do_cleanup().await { - if cfg!(any(test, feature = "test")) { - panic!("Failed to clean up vector index creator, err: {err:?}"); - } else { - warn!(err; "Failed to clean up vector index creator"); - } - } - return Err(update_err); - } - - Ok(()) - } - - /// Internal update implementation. - async fn do_update(&mut self, batch: &mut Batch) -> Result<()> { - let mut guard = self.stats.record_update(); - let n = batch.num_rows(); - guard.inc_row_count(n); - - for (col_id, creator) in &mut self.creators { - let Some(values) = batch.field_col_value(*col_id) else { - creator.push_nulls(n).map_err(build_err)?; - continue; - }; - - for i in 0..n { - let value = values.data.get_ref(i); - if value.is_null() { - creator.push_null().map_err(build_err)?; - } else if let ValueRef::Binary(bytes) = value { - let floats = bytes_to_f32_slice(bytes); - if floats.len() != creator.config().dim { - return VectorIndexBuildSnafu { - reason: format!( - "Vector dimension mismatch: expected {}, got {}", - creator.config().dim, - floats.len() - ), - } - .fail(); - } - creator.push_vector(&floats).map_err(build_err)?; - } else { - common_telemetry::debug!( - "Unexpected non-binary value for vector index column, treating as null, data_type: {:?}", - value.data_type() - ); - creator.push_null().map_err(build_err)?; - } - } - - // Check memory limit - abort index creation if exceeded - if let Some(threshold) = self.memory_usage_threshold { - let current_usage = creator.memory_usage(); - if current_usage > threshold { - warn!( - "Vector index memory usage {} exceeds threshold {}, aborting index creation, region_id: {}", - current_usage, threshold, self.metadata.region_id - ); - return VectorIndexBuildSnafu { - reason: format!( - "Memory usage {} exceeds threshold {}", - current_usage, threshold - ), - } - .fail(); - } - } - } - - Ok(()) - } - - /// Internal flat update implementation. - async fn do_update_flat(&mut self, batch: &RecordBatch) -> Result<()> { - let mut guard = self.stats.record_update(); - let n = batch.num_rows(); - guard.inc_row_count(n); - - for (col_id, creator) in &mut self.creators { - // This should never happen: creator exists but column not in metadata - let column_meta = self.metadata.column_by_id(*col_id).ok_or_else(|| { - VectorIndexBuildSnafu { - reason: format!( - "Column {} not found in region metadata, this is a bug", - col_id - ), - } - .build() - })?; - - let column_name = &column_meta.column_schema.name; - // Column not in batch is normal for flat format - treat as NULLs - let Some(column_array) = batch.column_by_name(column_name) else { - creator.push_nulls(n).map_err(build_err)?; - continue; - }; - - let binary_array = column_array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - VectorIndexBuildSnafu { - reason: format!( - "Column {} is not a binary array, got {:?}", - column_name, - column_array.data_type() - ), - } - .build() - })?; - - for i in 0..n { - if !binary_array.is_valid(i) { - creator.push_null().map_err(build_err)?; - } else { - let bytes = binary_array.value(i); - let floats = bytes_to_f32_slice(bytes); - if floats.len() != creator.config().dim { - return VectorIndexBuildSnafu { - reason: format!( - "Vector dimension mismatch: expected {}, got {}", - creator.config().dim, - floats.len() - ), - } - .fail(); - } - creator.push_vector(&floats).map_err(build_err)?; - } - } - - if let Some(threshold) = self.memory_usage_threshold { - let current_usage = creator.memory_usage(); - if current_usage > threshold { - warn!( - "Vector index memory usage {} exceeds threshold {}, aborting index creation, region_id: {}", - current_usage, threshold, self.metadata.region_id - ); - return VectorIndexBuildSnafu { - reason: format!( - "Memory usage {} exceeds threshold {}", - current_usage, threshold - ), - } - .fail(); - } - } - } - - Ok(()) - } - - /// Finishes index creation and writes to puffin. - /// Returns the number of rows and bytes written. - pub async fn finish( - &mut self, - puffin_writer: &mut SstPuffinWriter, - ) -> Result<(RowCount, ByteCount)> { - ensure!(!self.aborted, OperateAbortedIndexSnafu); - - if self.stats.row_count() == 0 { - // No IO is performed, no garbage to clean up - return Ok((0, 0)); - } - - let finish_res = self.do_finish(puffin_writer).await; - // Clean up garbage no matter finish successfully or not - if let Err(err) = self.do_cleanup().await { - if cfg!(any(test, feature = "test")) { - panic!("Failed to clean up vector index creator, err: {err:?}"); - } else { - warn!(err; "Failed to clean up vector index creator"); - } - } - - if finish_res.is_ok() { - INDEX_CREATE_ROWS_TOTAL - .with_label_values(&[TYPE_VECTOR_INDEX]) - .inc_by(self.stats.row_count() as u64); - INDEX_CREATE_BYTES_TOTAL - .with_label_values(&[TYPE_VECTOR_INDEX]) - .inc_by(self.stats.byte_count()); - } - - finish_res.map(|_| (self.stats.row_count(), self.stats.byte_count())) - } - - /// Internal finish implementation. - async fn do_finish(&mut self, puffin_writer: &mut SstPuffinWriter) -> Result<()> { - let mut guard = self.stats.record_finish(); - - for (id, creator) in &mut self.creators { - if creator.size() == 0 { - // No vectors to index - continue; - } - - let written_bytes = Self::do_finish_single_creator(*id, creator, puffin_writer).await?; - guard.inc_byte_count(written_bytes); - } - - Ok(()) - } - - /// Finishes a single column's vector index using the VectorIndexCreator trait. - async fn do_finish_single_creator( - col_id: ColumnId, - creator: &mut HnswVectorIndexCreator, - puffin_writer: &mut SstPuffinWriter, - ) -> Result { - let blob_name = format!("{}-{}", INDEX_BLOB_TYPE, col_id); - - let written_bytes = creator - .finish(puffin_writer, &blob_name, PutOptions::default()) - .await - .map_err(|e| { - VectorIndexFinishSnafu { - reason: e.to_string(), - } - .build() - })?; - - Ok(written_bytes) - } - - /// Aborts index creation and cleans up garbage. - pub async fn abort(&mut self) -> Result<()> { - if self.aborted { - return Ok(()); - } - self.aborted = true; - self.do_cleanup().await - } - - /// Cleans up temporary files. - async fn do_cleanup(&mut self) -> Result<()> { - let mut _guard = self.stats.record_cleanup(); - self.creators.clear(); - self.temp_file_provider.cleanup().await - } - - /// Returns the memory usage of the indexer. - pub fn memory_usage(&self) -> usize { - self.creators.values().map(|c| c.memory_usage()).sum() - } - - /// Returns the column IDs being indexed. - pub fn column_ids(&self) -> impl Iterator + '_ { - self.creators.keys().copied() - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Arc; - - use api::v1::SemanticType; - use datatypes::data_type::ConcreteDataType; - use datatypes::schema::ColumnSchema; - use datatypes::value::ValueRef; - use datatypes::vectors::{TimestampMillisecondVector, UInt8Vector, UInt64Vector}; - use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; - use store_api::storage::{ColumnId, FileId, RegionId}; - - use super::*; - use crate::read::BatchColumn; - - fn mock_region_metadata_with_vector() -> RegionMetadataRef { - let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 1, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("tag", ConcreteDataType::int64_datatype(), true), - semantic_type: SemanticType::Tag, - column_id: 2, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("vec", ConcreteDataType::vector_datatype(2), true), - semantic_type: SemanticType::Field, - column_id: 3, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "field_u64", - ConcreteDataType::uint64_datatype(), - true, - ), - semantic_type: SemanticType::Field, - column_id: 4, - }) - .primary_key(vec![2]); - - Arc::new(builder.build().unwrap()) - } - - fn new_batch_missing_vector_column(column_id: ColumnId, rows: usize) -> Batch { - let fields = vec![(0, SortField::new(ConcreteDataType::int64_datatype()))]; - let codec = DensePrimaryKeyCodec::with_fields(fields); - let primary_key = codec.encode([ValueRef::Int64(1)].into_iter()).unwrap(); - - let field = BatchColumn { - column_id, - data: Arc::new(UInt64Vector::from_iter_values(0..rows as u64)), - }; - - Batch::new( - primary_key, - Arc::new(TimestampMillisecondVector::from_values( - (0..rows).map(|i| i as i64).collect::>(), - )), - Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(0, rows))), - Arc::new(UInt8Vector::from_iter_values(std::iter::repeat_n(1, rows))), - vec![field], - ) - .unwrap() - } - - #[tokio::test] - async fn test_vector_index_missing_column_as_nulls() { - let tempdir = common_test_util::temp_dir::create_temp_dir( - "test_vector_index_missing_column_as_nulls_", - ); - let intm_mgr = IntermediateManager::init_fs(tempdir.path().to_string_lossy()) - .await - .unwrap(); - let region_metadata = mock_region_metadata_with_vector(); - - let mut vector_index_options = HashMap::new(); - vector_index_options.insert(3, VectorIndexOptions::default()); - - let mut indexer = VectorIndexer::new( - FileId::random(), - ®ion_metadata, - intm_mgr, - None, - &vector_index_options, - ) - .unwrap() - .unwrap(); - - let mut batch = new_batch_missing_vector_column(4, 3); - indexer.update(&mut batch).await.unwrap(); - - let creator = indexer.creators.get(&3).unwrap(); - assert_eq!(creator.size(), 0); - assert_eq!(creator.current_row_offset(), 3); - assert_eq!(creator.null_bitmap().len(), 3); - for idx in 0..3 { - assert!(creator.null_bitmap().contains(idx as u32)); - } - } -} diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index b49d2ba1e18..9404d4484fc 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -77,6 +77,18 @@ pub(crate) struct Json2TargetLayout { /// batching without changing the row group layout of newly written SSTs. pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; +/// Truncation length of min/max values in the parquet column index. +/// +/// Mito never reads the column index (it is skipped or stripped on load), so truncating +/// it is safe. Chunk statistics stay untruncated because pruning decodes primary keys +/// from them. Without truncation every page of a large string column stores its whole +/// min and max values, uncompressed. +/// +/// Truncated values are only bounds: code that starts reading the column index must not +/// decode primary keys from it. +pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option = + parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; + /// Applies the configured encoding to direct floating-point field columns. pub(crate) fn apply_float_field_encoding( mut builder: WriterPropertiesBuilder, @@ -901,6 +913,82 @@ mod tests { .await; } + #[tokio::test] + async fn test_column_index_truncates_large_field_values() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let handle = sst_file_handle(0, 1000); + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let metadata = build_test_binary_test_region_metadata(); + // A tag long enough that the encoded primary key exceeds the 64-byte truncation length. + let tag = "t".repeat(100); + let values: Vec> = (0..64) + .map(|i| format!("{i:08}").into_bytes().repeat(4 * 1024)) + .collect(); + let batch = new_record_batch_with_binary_values(&tag, &values); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + writer + .write_all_flat_as_primary_key( + new_flat_source_from_record_batches(vec![batch]), + None, + &WriteOptions::default(), + ) + .await + .unwrap(); + + let path = handle.file_path(FILE_DIR, PathType::Bare); + let bytes = object_store.read(&path).await.unwrap().to_bytes(); + let options = parquet::arrow::arrow_reader::ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options).unwrap(); + let parquet_meta = builder.metadata().clone(); + let schema = parquet_meta.file_metadata().schema_descr(); + let column = |name: &str| { + (0..schema.num_columns()) + .find(|i| schema.column(*i).name() == name) + .unwrap() + }; + + let field_index = &parquet_meta.column_index().unwrap()[0][column("field_0")]; + let parquet::file::page_index::column_index::ColumnIndexMetaData::BYTE_ARRAY(field_index) = + field_index + else { + panic!("unexpected column index: {field_index:?}"); + }; + let mut pages = 0; + for (min, max) in field_index + .min_values_iter() + .zip(field_index.max_values_iter()) + { + assert!(min.unwrap().len() <= 64); + assert!(max.unwrap().len() <= 64); + pages += 1; + } + assert!(pages > 1, "values should span several pages"); + + // Chunk statistics keep exact primary keys; pruning decodes them. + let pk_stats = parquet_meta + .row_group(0) + .column(column(store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME)) + .statistics() + .unwrap(); + let pk = new_primary_key(&[&tag]); + assert!(pk.len() > 64); + assert_eq!(pk_stats.max_bytes_opt().unwrap(), pk.as_slice()); + } + #[rstest::rstest] #[tokio::test] async fn test_write_multiple_files(#[values(1024, 4096)] write_buffer_size: usize) { @@ -1175,8 +1263,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut metrics = Metrics::new(WriteType::Flush); @@ -1477,6 +1563,32 @@ mod tests { assert!(cached.contains_row_group(3)); } + fn new_record_batch_with_binary_values(tag: &str, values: &[Vec]) -> RecordBatch { + let metadata = build_test_binary_test_region_metadata(); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + let num_rows = values.len(); + let mut tag_0_builder = StringDictionaryBuilder::::new(); + let mut pk_builder = BinaryDictionaryBuilder::::new(); + let pk = new_primary_key(&[tag]); + for _ in 0..num_rows { + tag_0_builder.append_value(tag); + pk_builder.append(&pk).unwrap(); + } + let columns: Vec = vec![ + Arc::new(tag_0_builder.finish()), + Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values( + values.iter().map(|v| v.as_slice()), + )), + Arc::new(TimestampMillisecondArray::from_iter_values( + 0..num_rows as i64, + )), + Arc::new(pk_builder.finish()), + Arc::new(UInt64Array::from_value(1000, num_rows)), + Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)), + ]; + RecordBatch::try_new(flat_schema, columns).unwrap() + } + fn new_record_batch_with_binary(tags: &[&str], start: usize, end: usize) -> RecordBatch { assert!(end >= start); let metadata = build_test_binary_test_region_metadata(); @@ -1695,8 +1807,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } } @@ -1816,8 +1926,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut metrics = Metrics::new(WriteType::Flush); diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 60b6a955e93..1c15afd386d 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -215,7 +215,7 @@ impl FileRange { selector, Some(TimeSeriesRowSelector::LastRow { after_merge: false }) ) { - // Only use LastRowReader if row group does not contain DELETE, all + // Only use the last row reader if row group does not contain DELETE, all // rows are selected, and filters that still run after this reader // cannot change which row is last. Tag filters are safe because a // tag is constant within a series. Timestamp and field filters are diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index a99fc201baf..3abba4d2db6 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -37,7 +37,7 @@ use datatypes::arrow::array::{ Array, ArrayRef, BinaryArray, DictionaryArray, UInt32Array, UInt64Array, }; use datatypes::arrow::compute::kernels::take::take; -use datatypes::arrow::datatypes::{Schema, SchemaRef}; +use datatypes::arrow::datatypes::{DataType as ArrowDataType, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::value::ValueRef; @@ -46,7 +46,7 @@ use mito_codec::row_converter::sparse::{ RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparsePrimaryKeyView, }; use mito_codec::row_converter::{ - CompositeValues, PrimaryKeyCodec, SparseOffsetsCache, build_primary_key_codec, + DensePrimaryKeyCodec, PrimaryKeyCodec, SparseOffsetsCache, build_primary_key_codec, }; use parquet::file::metadata::RowGroupMetaData; use snafu::{OptionExt, ResultExt, ensure}; @@ -407,11 +407,42 @@ impl FlatReadFormat { }; // First, apply flat format conversion. - let batch = match &self.parquet_adapter { + let mut batch = match &self.parquet_adapter { ParquetAdapter::Flat(_) => record_batch, ParquetAdapter::PrimaryKeyToFlat(p) => p.convert_batch(record_batch)?, }; + // Normalize nested field names and metadata to the SST's region metadata + // before schema compatibility and merging with memtables. This removes + // Parquet-added field IDs; equals_datatype also permits nested name differences. + for index in 0..batch.num_columns() { + let array = batch.column(index); + if !matches!(array.data_type(), ArrowDataType::Struct(_)) { + continue; + } + let field = batch.schema_ref().field(index); + let Some(column) = self.metadata().column_by_name(field.name()) else { + continue; + }; + let target = column.column_schema.data_type.as_arrow_type(); + if array.data_type() != &target && array.data_type().equals_datatype(&target) { + let array = + datatypes::arrow::compute::cast(array, &target).context(ComputeArrowSnafu)?; + let mut fields = batch.schema().fields().to_vec(); + fields[index] = Arc::new(field.clone().with_data_type(target)); + let mut columns = batch.columns().to_vec(); + columns[index] = array; + batch = RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + fields, + batch.schema().metadata().clone(), + )), + columns, + ) + .context(NewRecordBatchSnafu)?; + } + } + // Then apply sequence override if provided let Some(override_array) = override_sequence_array else { return Ok(batch); @@ -676,8 +707,9 @@ pub(crate) fn sst_column_id_indices(metadata: &RegionMetadata) -> HashMap { - let mut decoded_pk_values = Vec::with_capacity(distinct_keys.len()); - for key in &distinct_keys { - let pk_bytes = pk_values_array.value(*key as usize); - let decoded_value = codec.decode(pk_bytes).context(DecodeSnafu)?; - decoded_pk_values.push(decoded_value); - } - DecodedKeysInner::Dense(decoded_pk_values) - } + PrimaryKeyEncoding::Dense => DecodedKeysInner::Dense { + codec: codec + .as_dense() + .context(InvalidRecordBatchSnafu { + reason: "expected dense primary key codec", + })? + .clone(), + offsets: Vec::new(), + values: pk_values_array.clone(), + distinct_keys, + }, }; Ok(DecodedPrimaryKeys { @@ -748,10 +782,16 @@ pub fn decode_primary_keys( }) } -/// Eagerly decoded dense primary keys or lazily extracted sparse primary keys. +/// Encoded primary keys and encoding-specific lookup state. enum DecodedKeysInner { - /// Dense primary keys are decoded once and shared by all tag columns. - Dense(Vec), + /// Dense keys share positional offsets across projected tag columns. + Dense { + codec: DensePrimaryKeyCodec, + values: BinaryArray, + distinct_keys: Vec, + /// Initialized by positional access; full extraction needs no offsets. + offsets: Vec>, + }, /// Sparse primary keys stay encoded; each tag column extracts only its own /// values from the raw keys. Sparse { @@ -811,7 +851,7 @@ impl DecodedPrimaryKeys { /// Gets a tag column array by column id and data type. /// /// For sparse encoding, extracts the column lazily from the encoded keys. - /// For dense encoding, uses pk_index to get values from the decoded keys. + /// For dense encoding, uses pk_index to decode only the requested field. pub fn get_tag_column( &mut self, column_id: ColumnId, @@ -827,18 +867,25 @@ impl DecodedPrimaryKeys { // Gets values from the primary key. let values_vector = match inner { - DecodedKeysInner::Dense(decoded_pk_values) => { - let mut builder = column_type.create_mutable_vector(decoded_pk_values.len()); - for decoded in decoded_pk_values { - let CompositeValues::Dense(dense) = decoded else { - return InvalidRecordBatchSnafu { - reason: "expected dense primary key values", - } - .fail(); - }; - let pk_idx = pk_index.expect("pk_index required for dense encoding"); - if pk_idx < dense.len() { - builder.push_value_ref(&dense[pk_idx].1.as_value_ref()); + DecodedKeysInner::Dense { + codec, + values, + distinct_keys, + offsets, + } => { + let pk_idx = pk_index.context(InvalidRecordBatchSnafu { + reason: "pk_index required for dense encoding", + })?; + let mut builder = column_type.create_mutable_vector(distinct_keys.len()); + if offsets.is_empty() { + *offsets = vec![Vec::new(); distinct_keys.len()]; + } + for (&key, offsets) in distinct_keys.iter().zip(offsets) { + if pk_idx < codec.num_fields() { + let value = codec + .decode_value_at(values.value(key as usize), pk_idx, offsets) + .context(DecodeSnafu)?; + builder.push_value_ref(&value.as_value_ref()); } else { builder.push_null(); } @@ -872,6 +919,49 @@ impl DecodedPrimaryKeys { } } + /// Materializes all Dense tags in source-schema order with one sequential + /// traversal per key. Values are consumed immediately by their Arrow builders. + pub fn get_dense_tag_columns(&self) -> Result> { + let DecodedKeysInner::Dense { + codec, + values, + distinct_keys, + .. + } = &self.inner + else { + return InvalidRecordBatchSnafu { + reason: "expected dense primary key values", + } + .fail(); + }; + let mut builders: Vec<_> = codec + .fields() + .iter() + .map(|(_, field)| field.data_type().create_mutable_vector(distinct_keys.len())) + .collect(); + let mut value_buf = Vec::new(); + for &key in distinct_keys { + codec + .decode_dense_with(values.value(key as usize), &mut value_buf, |pos, value| { + builders[pos].push_value_ref(&value); + }) + .context(DecodeSnafu)?; + } + codec + .fields() + .iter() + .zip(builders) + .map(|((_, field), mut builder)| { + let values = builder.to_vector().to_arrow_array(); + if field.data_type().is_string() { + Ok(Arc::new(DictionaryArray::new(self.keys_array.clone(), values)) as ArrayRef) + } else { + take(&values, &self.keys_array, None).context(ComputeArrowSnafu) + } + }) + .collect() + } + /// Gets multiple sparse tag column arrays in one pass over the distinct keys, /// sharing each key's offset discovery between the columns. /// @@ -1004,6 +1094,8 @@ impl FlatConvertFormat { }) .collect(); decoded_columns.extend(decoded_pks.get_sparse_tag_columns(&columns)?); + } else if self.projected_primary_keys.len() == self.metadata.primary_key.len() { + decoded_columns.extend(decoded_pks.get_dense_tag_columns()?); } else { for (column_id, pk_index, column_index) in &self.projected_primary_keys { let column_metadata = &self.metadata.column_metadatas[*column_index]; @@ -1086,6 +1178,138 @@ mod tests { flat_sst_arrow_schema_column_num, override_pk_field_to_binary, to_flat_sst_arrow_schema, }; + #[test] + fn dense_tag_columns_match_eager_decoding() { + use datatypes::arrow::datatypes::UInt32Type; + use datatypes::value::Value; + use datatypes::vectors::Helper; + use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; + + let columns = [ + (17, ConcreteDataType::string_datatype()), + (9, ConcreteDataType::int64_datatype()), + (2, ConcreteDataType::binary_datatype()), + ]; + let codec = DensePrimaryKeyCodec::with_fields( + columns + .iter() + .map(|(id, ty)| (*id, SortField::new(ty.clone()))) + .collect(), + ); + let rows = [ + vec![ + Value::from("中文\0abcdefgh"), + Value::Int64(-42), + Value::Binary(vec![0, 1, 255].into()), + ], + vec![Value::from(""), Value::Null, Value::Binary(vec![].into())], + vec![Value::Null, Value::Int64(i64::MAX), Value::Null], + ]; + let mut encoded: Vec<_> = rows + .iter() + .map(|row| codec.encode(row.iter().map(Value::as_value_ref)).unwrap()) + .collect(); + // Unreferenced dictionary entries must not be inspected. + encoded.push(vec![255]); + let row_keys = [2, 2, 0, 1, 0, 0]; + let pk = DictionaryArray::::new( + UInt32Array::from(row_keys.to_vec()), + Arc::new(BinaryArray::from_iter_values(&encoded)), + ); + let batch = RecordBatch::try_from_iter([ + ( + "ts", + Arc::new(TimestampMillisecondArray::from_iter_values(0..6)) as ArrayRef, + ), + ("__primary_key", Arc::new(pk) as ArrayRef), + ( + "__sequence", + Arc::new(UInt64Array::from(vec![1; 6])) as ArrayRef, + ), + ( + "__op_type", + Arc::new(UInt8Array::from(vec![1; 6])) as ArrayRef, + ), + ]) + .unwrap(); + let mut metadata = RegionMetadataBuilder::new(RegionId::new(1, 1)); + for (id, ty) in columns.iter().rev() { + metadata.push_column_metadata(ColumnMetadata { + column_id: *id, + column_schema: ColumnSchema::new(format!("tag_{id}"), ty.clone(), true), + semantic_type: SemanticType::Tag, + }); + } + metadata + .push_column_metadata(ColumnMetadata { + column_id: 23, + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + }) + .primary_key(vec![17, 9, 2]); + let metadata = Arc::new(metadata.build().unwrap()); + let format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new([17, 9, 2, 23]), + Some(batch.schema()), + "test", + false, + ) + .unwrap(); + for (offset, len) in [(0, 6), (1, 4), (3, 0)] { + let mut decoded = decode_primary_keys(&codec, &batch.slice(offset, len)).unwrap(); + let all = decoded.get_dense_tag_columns().unwrap(); + let converted = format + .convert_batch(batch.slice(offset, len), None) + .unwrap(); + // Out-of-order and repeated projections share the same key offsets. + for pos in [2, 0, 1, 0] { + let array = decoded + .get_tag_column(columns[pos].0, Some(pos), &columns[pos].1) + .unwrap(); + assert_eq!(array.to_data(), all[pos].to_data()); + assert_eq!(converted.column(pos).to_data(), all[pos].to_data()); + let actual = Helper::try_into_vector(array).unwrap(); + for (row, &key) in row_keys[offset..offset + len].iter().enumerate() { + let eager = codec + .decode_dense_without_column_id(&encoded[key as usize]) + .unwrap(); + assert_eq!(actual.get(row), eager[pos]); + } + } + // A position beyond the source schema remains NULL, never decoded + // with a newer schema's field order or type. + let missing = decoded + .get_tag_column(99, Some(3), &ConcreteDataType::int64_datatype()) + .unwrap(); + assert_eq!(missing.null_count(), len); + assert!(decoded.get_tag_column(17, None, &columns[0].1).is_err()); + } + let mut malformed_columns = batch.columns().to_vec(); + malformed_columns[1] = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0; 6]), + Arc::new(BinaryArray::from(vec![&[0, 1][..]])), + )); + let malformed = RecordBatch::try_new(batch.schema(), malformed_columns).unwrap(); + assert!(format.convert_batch(malformed.clone(), None).is_err()); + // A partial projection must still avoid decoding an unneeded broken suffix. + let partial = FlatReadFormat::new( + metadata, + ReadColumns::new([17, 23]), + Some(batch.schema()), + "test", + false, + ) + .unwrap(); + let partial = partial.convert_batch(malformed, None).unwrap(); + let tag = Helper::try_into_vector(partial.column(0).clone()).unwrap(); + assert!((0..6).all(|row| tag.get(row).is_null())); + } + /// Builds a `RegionMetadata` with the given number of tags and fields. fn build_metadata( num_tags: usize, diff --git a/src/mito2/src/sst/parquet/json_align.rs b/src/mito2/src/sst/parquet/json_align.rs index 8af88f6ac0f..93ecffa072f 100644 --- a/src/mito2/src/sst/parquet/json_align.rs +++ b/src/mito2/src/sst/parquet/json_align.rs @@ -22,7 +22,7 @@ use datatypes::arrow::array::{ArrayRef, new_null_array}; use datatypes::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::{JsonMetadata, is_json2_extension_type}; -use datatypes::json::JsonSettings; +use datatypes::json::{JsonSettings, TypeHintMismatchPolicy}; use datatypes::vectors::json::array::JsonArray; use datatypes::vectors::json::json2_physical_data_type; use futures::Stream; @@ -251,10 +251,11 @@ fn rewrite_array( settings: &RewriteSettings, ) -> Result { JsonArray::from(source_array) - .rewrite_to_v2( + .rewrite_to_v2_with_type_hint_mismatch_policy( source_field, &settings.logical_settings, &settings.target_layout, + TypeHintMismatchPolicy::CoerceOrNull, ) .context(DataTypeMismatchSnafu) } @@ -330,7 +331,9 @@ mod tests { Array, ArrayRef, BinaryArray, Int64Array, StringArray, StringViewArray, StructArray, }; use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema}; - use datatypes::extension::json::Json2ExtensionType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonTypeHint; + use datatypes::prelude::ConcreteDataType; use datatypes::types::parse_string_to_jsonb; use futures::{StreamExt, stream}; @@ -783,6 +786,75 @@ mod tests { assert_eq!(int_array([10, 20]).as_ref(), output.column(3).as_ref()); } + #[tokio::test] + async fn test_rewrite_keeps_rows_with_invalid_json2_settings() { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + inverted_index: false, + }], + Some(0), + ) + .unwrap(); + let target_type = json2_physical_data_type(&settings); + let output_schema = schema([ + Field::new("j", target_type.clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings.clone()))), + ), + Field::new("value", DataType::Int64, true), + ]); + let source = Arc::new(BinaryArray::from_iter([ + Some(parse_string_to_jsonb(r#"{"kind":"valid"}"#).unwrap()), + Some(parse_string_to_jsonb(r#"{"kind":1}"#).unwrap()), + ])) as ArrayRef; + let input = RecordBatch::try_new( + schema([ + Field::new("j", DataType::Binary, true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("value", DataType::Int64, true), + ]), + vec![source, int_array([10, 20])], + ) + .unwrap(); + let columns = HashMap::from([( + "j".to_string(), + Json2TargetLayout { + extension_metadata: serde_json::to_string(&JsonMetadata::new(settings.clone())) + .unwrap(), + target_layout: settings, + }, + )]); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![true, true], + output_schema, + AlignMode::Rewrite { columns }, + ) + .unwrap(); + + let output = aligner.next().await.unwrap().unwrap(); + assert_eq!(2, output.num_rows()); + assert_eq!( + 10, + output + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + ); + assert_eq!( + 20, + output + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(1) + ); + } + #[test] fn test_rewrite_rejects_mismatched_output_layout() { let columns = HashMap::from([( diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index 4362440e415..d5401f4a6de 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -341,8 +341,7 @@ pub(crate) fn build_bulk_filter_plan( predicate: Option<&Predicate>, ) -> BulkFilterPlan { let metadata = read_format.metadata(); - // Bulk memtable only needs simple binary filters here. Any filter that - // cannot be reduced to a SimpleFilterContext stays out of this fast path. + // Filters not recognized by SimpleFilterContext stay out of this fast path. let simple_filters: Vec = predicate .into_iter() .flat_map(|predicate| { diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index f54812f97d4..a3cbdef90f1 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -14,8 +14,6 @@ //! Parquet reader. -#[cfg(feature = "vector_index")] -use std::collections::BTreeSet; use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -55,8 +53,6 @@ use table::predicate::Predicate; use crate::cache::index::result_cache::PredicateKey; use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta}; -#[cfg(feature = "vector_index")] -use crate::error::ApplyVectorIndexSnafu; use crate::error::{ ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu, UnexpectedSnafu, @@ -79,8 +75,6 @@ use crate::sst::index::fulltext_index::applier::{ use crate::sst::index::inverted_index::applier::{ InvertedIndexApplierRef, InvertedIndexApplyMetrics, }; -#[cfg(feature = "vector_index")] -use crate::sst::index::vector_index::applier::VectorIndexApplierRef; use crate::sst::parquet::file_range::{ FileRangeContext, FileRangeContextRef, PartitionFilterContext, PreFilterMode, RangeBase, }; @@ -167,7 +161,6 @@ fn should_read_pk_as_binary_with_limit( } const INDEX_TYPE_INVERTED: &str = "inverted"; const INDEX_TYPE_BLOOM: &str = "bloom filter"; -const INDEX_TYPE_VECTOR: &str = "vector"; macro_rules! handle_index_error { ($err:expr, $file_handle:expr, $index_type:expr) => { @@ -215,12 +208,6 @@ pub struct ParquetReaderBuilder { inverted_index_appliers: [Option; 2], bloom_filter_index_appliers: [Option; 2], fulltext_index_appliers: [Option; 2], - /// Vector index applier for KNN search. - #[cfg(feature = "vector_index")] - vector_index_applier: Option, - /// Over-fetched k for vector index scan. - #[cfg(feature = "vector_index")] - vector_index_k: Option, /// Expected metadata of the region while reading the SST. /// This is usually the latest metadata of the region. The reader use /// it get the correct column id of a column by name. @@ -262,10 +249,6 @@ impl ParquetReaderBuilder { inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], fulltext_index_appliers: [None, None], - #[cfg(feature = "vector_index")] - vector_index_applier: None, - #[cfg(feature = "vector_index")] - vector_index_k: None, expected_metadata: None, compaction: false, pre_filter_mode: PreFilterMode::All, @@ -352,19 +335,6 @@ impl ParquetReaderBuilder { self } - /// Attaches the vector index applier to the builder. - #[cfg(feature = "vector_index")] - #[must_use] - pub(crate) fn vector_index_applier( - mut self, - applier: Option, - k: Option, - ) -> Self { - self.vector_index_applier = applier; - self.vector_index_k = k; - self - } - /// Attaches the expected metadata to the builder. #[must_use] pub fn expected_metadata(mut self, expected_metadata: Option) -> Self { @@ -713,7 +683,7 @@ impl ParquetReaderBuilder { && context .version .range_indexes - .contains(&self.file_handle.file_id().file_id())) + .contains_key(&self.file_handle.file_id().file_id())) .then(|| context.store.clone()) }); let context = FileRangeContext::new( @@ -851,7 +821,14 @@ impl ParquetReaderBuilder { let metadata = metadata_loader.load(cache_metrics).await?; let decoded = if self.cache_strategy.sst_meta_cache_enabled() { - let metadata = prepare_sst_meta(file_path, metadata, None, page_index_policy).await?; + let metadata = prepare_sst_meta( + file_path, + metadata, + None, + page_index_policy, + &self.cache_strategy.sst_meta_runtime(), + ) + .await?; let decoded = metadata.decoded(); match metadata { SstMetaPreparation::Prepared(metadata) => { @@ -974,19 +951,6 @@ impl ParquetReaderBuilder { ) .await; } - #[cfg(feature = "vector_index")] - { - self.prune_row_groups_by_vector_index( - row_group_size, - num_row_groups, - &mut output, - metrics, - ) - .await; - if output.is_empty() { - return output; - } - } output } @@ -1252,49 +1216,6 @@ impl ParquetReaderBuilder { pruned } - /// Prunes row groups by vector index results. - #[cfg(feature = "vector_index")] - async fn prune_row_groups_by_vector_index( - &self, - row_group_size: usize, - num_row_groups: usize, - output: &mut RowGroupSelection, - metrics: &mut ReaderFilterMetrics, - ) { - let Some(applier) = &self.vector_index_applier else { - return; - }; - let Some(k) = self.vector_index_k else { - return; - }; - if !self.file_handle.meta_ref().vector_index_available() { - return; - } - - let file_size_hint = self.file_handle.meta_ref().index_file_size(); - let apply_res = applier - .apply_with_k(self.file_handle.index_id(), Some(file_size_hint), k) - .await; - let row_ids = match apply_res { - Ok(res) => res.row_offsets, - Err(err) => { - handle_index_error!(err, self.file_handle, INDEX_TYPE_VECTOR); - return; - } - }; - - let selection = match vector_selection_from_offsets(row_ids, row_group_size, num_row_groups) - { - Ok(selection) => selection, - Err(err) => { - handle_index_error!(err, self.file_handle, INDEX_TYPE_VECTOR); - return; - } - }; - metrics.rows_vector_selected += selection.row_count(); - apply_selection_and_update_metrics(output, &selection, metrics, INDEX_TYPE_VECTOR); - } - async fn prune_row_groups_by_fulltext_bloom( &self, row_group_size: usize, @@ -1518,29 +1439,6 @@ fn apply_selection_and_update_metrics( *output = intersection; } -#[cfg(feature = "vector_index")] -fn vector_selection_from_offsets( - row_offsets: Vec, - row_group_size: usize, - num_row_groups: usize, -) -> Result { - let mut row_ids = BTreeSet::new(); - for offset in row_offsets { - let row_id = u32::try_from(offset).map_err(|_| { - ApplyVectorIndexSnafu { - reason: format!("Row offset {} exceeds u32::MAX", offset), - } - .build() - })?; - row_ids.insert(row_id); - } - Ok(RowGroupSelection::from_row_ids( - row_ids, - row_group_size, - num_row_groups, - )) -} - fn all_required_row_groups_searched( required_row_groups: &RowGroupSelection, cached_row_groups: &RowGroupSelection, @@ -1566,8 +1464,6 @@ pub(crate) struct ReaderFilterMetrics { pub(crate) rg_minmax_filtered: usize, /// Number of row groups filtered by bloom filter index. pub(crate) rg_bloom_filtered: usize, - /// Number of row groups filtered by vector index. - pub(crate) rg_vector_filtered: usize, /// Number of rows in row group before filtering. pub(crate) rows_total: usize, @@ -1577,10 +1473,6 @@ pub(crate) struct ReaderFilterMetrics { pub(crate) rows_inverted_filtered: usize, /// Number of rows in row group filtered by bloom filter index. pub(crate) rows_bloom_filtered: usize, - /// Number of rows filtered by vector index. - pub(crate) rows_vector_filtered: usize, - /// Number of rows selected by vector index. - pub(crate) rows_vector_selected: usize, /// Number of rows filtered by precise filter. pub(crate) rows_precise_filtered: usize, @@ -1626,14 +1518,11 @@ impl ReaderFilterMetrics { self.rg_inverted_filtered += other.rg_inverted_filtered; self.rg_minmax_filtered += other.rg_minmax_filtered; self.rg_bloom_filtered += other.rg_bloom_filtered; - self.rg_vector_filtered += other.rg_vector_filtered; self.rows_total += other.rows_total; self.rows_fulltext_filtered += other.rows_fulltext_filtered; self.rows_inverted_filtered += other.rows_inverted_filtered; self.rows_bloom_filtered += other.rows_bloom_filtered; - self.rows_vector_filtered += other.rows_vector_filtered; - self.rows_vector_selected += other.rows_vector_selected; self.rows_precise_filtered += other.rows_precise_filtered; self.fulltext_index_cache_hit += other.fulltext_index_cache_hit; @@ -1685,9 +1574,6 @@ impl ReaderFilterMetrics { READ_ROW_GROUPS_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rg_bloom_filtered as u64); - READ_ROW_GROUPS_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rg_vector_filtered as u64); PRECISE_FILTER_ROWS_TOTAL .with_label_values(&["parquet"]) @@ -1704,9 +1590,6 @@ impl ReaderFilterMetrics { READ_ROWS_IN_ROW_GROUP_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rows_bloom_filtered as u64); - READ_ROWS_IN_ROW_GROUP_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rows_vector_filtered as u64); } fn update_index_metrics(&mut self, index_type: &str, row_group_count: usize, row_count: usize) { @@ -1723,67 +1606,11 @@ impl ReaderFilterMetrics { self.rg_bloom_filtered += row_group_count; self.rows_bloom_filtered += row_count; } - INDEX_TYPE_VECTOR => { - self.rg_vector_filtered += row_group_count; - self.rows_vector_filtered += row_count; - } _ => {} } } } -#[cfg(all(test, feature = "vector_index"))] -mod vector_index_tests { - use super::*; - - #[test] - fn test_vector_selection_from_offsets() { - let row_group_size = 4; - let num_row_groups = 3; - let selection = - vector_selection_from_offsets(vec![0, 1, 5, 9], row_group_size, num_row_groups) - .unwrap(); - - assert_eq!(selection.row_group_count(), 3); - assert_eq!(selection.row_count(), 4); - assert!(selection.contains_non_empty_row_group(0)); - assert!(selection.contains_non_empty_row_group(1)); - assert!(selection.contains_non_empty_row_group(2)); - } - - #[test] - fn test_vector_selection_from_offsets_out_of_range() { - let row_group_size = 4; - let num_row_groups = 2; - let selection = vector_selection_from_offsets( - vec![0, 7, u64::from(u32::MAX) + 1], - row_group_size, - num_row_groups, - ); - assert!(selection.is_err()); - } - - #[test] - fn test_vector_selection_updates_metrics() { - let row_group_size = 4; - let total_rows = 8; - let mut output = RowGroupSelection::new(row_group_size, total_rows); - let selection = vector_selection_from_offsets(vec![1], row_group_size, 2).unwrap(); - let mut metrics = ReaderFilterMetrics::default(); - - apply_selection_and_update_metrics( - &mut output, - &selection, - &mut metrics, - INDEX_TYPE_VECTOR, - ); - - assert_eq!(metrics.rg_vector_filtered, 1); - assert_eq!(metrics.rows_vector_filtered, 7); - assert_eq!(output.row_count(), 1); - } -} - /// Metrics for parquet metadata cache operations. #[derive(Default, Clone, Copy)] pub struct MetadataCacheMetrics { @@ -2187,7 +2014,12 @@ impl SimpleFilterContext { expected_meta: Option<&RegionMetadata>, expr: &Expr, ) -> Option { - let filter = SimpleFilterEvaluator::try_new(expr)?; + let filter = SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| { + expected_meta + .unwrap_or(sst_meta) + .column_by_name(name) + .map(|column| column.column_schema.data_type.as_arrow_type()) + })?; let expr_str = format!("{expr:?}"); let (column_metadata, maybe_filter) = match expected_meta { Some(meta) => { @@ -3332,6 +3164,15 @@ mod tests { )); } + #[test] + fn test_nullable_label_filter_uses_missing_sst_column_default() { + let metadata = Arc::new(sst_region_metadata()); + let expected = expected_metadata_with_reused_tag_name(&metadata); + let expr = col("tag_0").is_null().or(col("tag_0").not_eq(lit("lo"))); + let ctx = SimpleFilterContext::new_opt(&metadata, Some(&expected), &expr).unwrap(); + assert!(matches!(ctx.filter(), MaybeFilter::Matched)); + } + #[test] fn test_simple_filter_context_drops_mismatched_field_filter() { let (sst_metadata, latest_metadata) = mock_metadata(); diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index b9ef15add34..932654dfdd8 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -62,7 +62,8 @@ use crate::sst::parquet::flat_format::{ }; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{ - PARQUET_METADATA_KEY, SstInfo, WriteOptions, apply_float_field_encoding, + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, + apply_float_field_encoding, }; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, @@ -544,7 +545,7 @@ where .set_compression(Compression::ZSTD(ZstdLevel::default())) .set_encoding(Encoding::PLAIN) .set_max_row_group_row_count(Some(opts.row_group_size)) - .set_column_index_truncate_length(None) + .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); let ts_col = ColumnPath::new(vec![ self.metadata.time_index_column().column_schema.name.clone(), diff --git a/src/mito2/src/sst/primary_key.rs b/src/mito2/src/sst/primary_key.rs index 5fb258052c7..7752a652f80 100644 --- a/src/mito2/src/sst/primary_key.rs +++ b/src/mito2/src/sst/primary_key.rs @@ -270,12 +270,12 @@ mod tests { } fn version_with_files(metadata: RegionMetadataRef, files: Vec) -> VersionRef { - let mutable = Arc::new(TimePartitions::new( + let mutable = TimePartitions::new( metadata.clone(), Arc::new(TimeSeriesMemtableBuilder::default()), 0, None, - )); + ); Arc::new( VersionBuilder::new(metadata, mutable) .add_files(new_noop_file_purger(), files.into_iter()) diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index dbfbc53db25..d5104bb16b7 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -42,11 +42,13 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_telemetry::{debug, warn}; use common_test_util::temp_dir::{TempDir, create_temp_dir}; +use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; use datatypes::arrow::array::{TimestampMillisecondArray, UInt8Array, UInt64Array}; use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; +use log_store::ObjectStoreLogStore; use log_store::kafka::log_store::KafkaLogStore; use log_store::raft_engine::log_store::RaftEngineLogStore; use log_store::test_util::log_store_util; @@ -56,7 +58,7 @@ use object_store::layers::mock::{ Buffer, Deleter, Metadata, MockLayer, MockLayerBuilder, OpDelete, Result as MockResult, Writer, }; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; -use object_store::services::Fs; +use object_store::services::{Fs, Memory}; use rskafka::client::partition::{Compression, UnknownTopicHandling}; use rskafka::client::{Client, ClientBuilder}; use rskafka::record::Record; @@ -79,7 +81,7 @@ use crate::error::Result; use crate::flush::{WriteBufferManager, WriteBufferManagerRef}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::manifest::storage::{is_checkpoint_file, is_delta_file}; -use crate::read::{Batch, BatchBuilder, BatchReader}; +use crate::read::{Batch, BatchBuilder}; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::sst::FormatType; use crate::sst::file_purger::{FilePurgerRef, NoopFilePurger}; @@ -492,6 +494,49 @@ impl TestEnv { self.new_mito_engine_with_plugins(config, plugins).await } + /// Creates a new engine with specific config on `log_store` under this env. + pub(crate) async fn create_engine_with_log_store( + &mut self, + config: MitoConfig, + log_store: Arc, + ) -> MitoEngine { + let object_store_manager = Arc::new(self.create_object_store_manager()); + self.object_store_manager = Some(object_store_manager.clone()); + + let data_home = self.data_home().display().to_string(); + MitoEngine::new( + &data_home, + config, + log_store, + object_store_manager, + self.schema_metadata_manager.clone(), + self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), + Plugins::new(), + ) + .await + .unwrap() + } + + /// Creates a new engine with specific config on an [ObjectStoreLogStore] that + /// keeps its objects in memory, opened as datanode 0 in the standalone + /// generation. Returns the store so that a test can drive its testing hooks. + pub(crate) async fn create_engine_with_object_store_wal( + &mut self, + config: MitoConfig, + wal_config: &ObjectStoreWalConfig, + ) -> (MitoEngine, Arc) { + let wal_object_store = ObjectStore::new(Memory::default()).unwrap(); + let log_store = + ObjectStoreLogStore::try_new(wal_object_store, wal_config, 0, STANDALONE_GENERATION) + .await + .unwrap(); + let engine = self + .create_engine_with_log_store(config, log_store.clone()) + .await; + (engine, log_store) + } + /// Creates a new engine with specific config and existing logstore and object store manager. pub async fn create_follower_engine(&mut self, config: MitoConfig) -> MitoEngine { self.new_mito_engine(config).await @@ -1106,34 +1151,6 @@ pub(crate) fn ts_ms_value(data: i64) -> v1::Value { } } -/// A reader for test that pop [Batch] from a vector. -pub struct VecBatchReader { - batches: Vec, -} - -impl VecBatchReader { - pub fn new(batches: &[Batch]) -> VecBatchReader { - let batches = batches.iter().rev().cloned().collect(); - - VecBatchReader { batches } - } -} - -#[async_trait::async_trait] -impl BatchReader for VecBatchReader { - async fn next_batch(&mut self) -> Result> { - Ok(self.batches.pop()) - } -} - -impl Iterator for VecBatchReader { - type Item = Result; - - fn next(&mut self) -> Option> { - self.batches.pop().map(Ok) - } -} - pub fn new_batch_builder( primary_key: &[u8], timestamps: &[i64], @@ -1177,19 +1194,6 @@ pub fn new_batch( .unwrap() } -/// Ensure the reader returns batch as `expect`. -pub async fn check_reader_result(reader: &mut R, expect: &[Batch]) { - let mut result = Vec::new(); - while let Some(mut batch) = reader.next_batch().await.unwrap() { - batch.remove_pk_values(); - result.push(batch); - } - - assert_eq!(expect, result); - // Next call to `next_batch()` still returns None. - assert!(reader.next_batch().await.unwrap().is_none()); -} - /// A mock [WriteBufferManager] that supports controlling whether to flush/stall. #[derive(Debug, Default)] pub struct MockWriteBufferManager { diff --git a/src/mito2/src/test_util/bench_util.rs b/src/mito2/src/test_util/bench_util.rs index ac635d2799a..e0334676e5b 100644 --- a/src/mito2/src/test_util/bench_util.rs +++ b/src/mito2/src/test_util/bench_util.rs @@ -21,6 +21,7 @@ use api::v1::value::ValueData; use api::v1::{Row, Rows, SemanticType}; use datafusion_common::Column; use datafusion_expr::{Expr, lit}; +use datatypes::arrow::datatypes::SchemaRef; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; @@ -30,7 +31,7 @@ use rand::seq::IndexedRandom; use store_api::metadata::{ ColumnMetadata, RegionMetadata, RegionMetadataBuilder, RegionMetadataRef, }; -use store_api::storage::RegionId; +use store_api::storage::{ColumnId, RegionId}; use table::predicate::Predicate; use crate::memtable::KeyValues; @@ -42,6 +43,23 @@ use crate::sst::parquet::flat_format::FlatReadFormat; use crate::sst::parquet::reader::SimpleFilterContext; use crate::test_util::memtable_util::region_metadata_to_row_schema; +/// Builds the actual encoded-PK-to-flat conversion used by the SST reader. +pub fn pk_materializer_for_bench( + metadata: RegionMetadataRef, + columns: Vec, + file_schema: SchemaRef, +) -> impl Fn(RecordBatch) -> crate::error::Result { + let format = FlatReadFormat::new( + metadata, + ReadColumns::new(columns), + Some(file_schema), + "bench", + false, + ) + .unwrap(); + move |batch| format.convert_batch(batch, None) +} + /// Builds a precise-filter benchmark with tags left encoded in the primary key. pub fn tag_filter_for_bench( metadata: RegionMetadataRef, diff --git a/src/mito2/src/test_util/version_util.rs b/src/mito2/src/test_util/version_util.rs index 1c0d7bfebc8..44827b35cc1 100644 --- a/src/mito2/src/test_util/version_util.rs +++ b/src/mito2/src/test_util/version_util.rs @@ -132,12 +132,7 @@ impl VersionControlBuilder { pub(crate) fn build_version(&self) -> Version { let metadata = Arc::new(self.metadata.clone()); - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - self.memtable_builder.clone(), - 0, - None, - )); + let mutable = TimePartitions::new(metadata.clone(), self.memtable_builder.clone(), 0, None); VersionBuilder::new(metadata, mutable) .add_files(self.file_purger.clone(), self.files.values().cloned()) .build() diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index fa549622898..772aac336cf 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -36,7 +36,7 @@ use store_api::logstore::provider::Provider; use store_api::logstore::{AppendBatchResponse, LogStore, WalIndex}; use store_api::storage::RegionId; -use crate::error::{BuildEntrySnafu, DeleteWalSnafu, Result, WriteWalSnafu}; +use crate::error::{BuildEntrySnafu, DeleteWalSnafu, Result, WaitWalDurableSnafu, WriteWalSnafu}; use crate::wal::entry_reader::{LogStoreEntryReader, WalEntryReader}; use crate::wal::raw_entry_reader::{LogStoreRawEntryReader, RegionRawEntryReader}; @@ -45,6 +45,36 @@ pub type EntryId = store_api::logstore::entry::Id; /// A stream that yields tuple of WAL entry id and corresponding entry. pub type WalEntryStream<'a> = BoxStream<'a, Result<(EntryId, WalEntry)>>; +/// Waits until the WAL of one region is durable through an entry id. +/// +/// A flush calls it before it records an entry id as flushed in the manifest: +/// a log store that acknowledges an append before its entries are durable +/// hands out entry ids that a crash can lose, and a manifest watermark that +/// names such an id would skip the entries assigned again after a restart. +#[derive(Clone)] +pub(crate) struct DurabilityBarrier( + Arc BoxFuture<'static, Result<()>> + Send + Sync>, +); + +impl DurabilityBarrier { + /// Returns once the region is durable through `entry_id`. + pub(crate) async fn wait(&self, entry_id: EntryId) -> Result<()> { + (self.0)(entry_id).await + } + + /// A barrier that never waits, for tests without a log store. + #[cfg(test)] + pub(crate) fn noop() -> Self { + Self(Arc::new(|_| Box::pin(async { Ok(()) }))) + } +} + +impl std::fmt::Debug for DurabilityBarrier { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("DurabilityBarrier") + } +} + /// Write ahead log. /// /// All regions in the engine shares the same WAL instance. @@ -104,6 +134,30 @@ impl Wal { } } + /// Returns the [DurabilityBarrier] of the region written through `provider`. + pub(crate) fn durability_barrier( + &self, + region_id: RegionId, + provider: &Provider, + ) -> DurabilityBarrier { + let store = self.store.clone(); + let provider = provider.clone(); + DurabilityBarrier(Arc::new(move |entry_id| { + let store = store.clone(); + let provider = provider.clone(); + Box::pin(async move { + if let Provider::Noop = provider { + return Ok(()); + } + store + .wait_durable(&provider, entry_id) + .await + .map_err(BoxedError::new) + .context(WaitWalDurableSnafu { region_id }) + }) + })) + } + /// Returns a [WalEntryReader] pub(crate) fn wal_entry_reader( &self, diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 24eec5dc2b3..2eda4fff49f 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -35,7 +35,7 @@ mod handle_write; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; use common_base::Plugins; @@ -198,6 +198,7 @@ impl WorkerGroup { let index_build_job_pool = Arc::new(LocalScheduler::new(config.max_background_index_builds)); let series_index_store = series_index_store_from_config(&config, data_home).await?; + let series_index_disk_usage = Arc::new(AtomicU64::new(0)); 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)); @@ -251,6 +252,7 @@ impl WorkerGroup { write_buffer_manager: write_buffer_manager.clone(), index_build_job_pool: index_build_job_pool.clone(), series_index_store: series_index_store.clone(), + series_index_disk_usage: series_index_disk_usage.clone(), flush_job_pool: flush_job_pool.clone(), compact_job_pool: compact_job_pool.clone(), purge_scheduler: purge_scheduler.clone(), @@ -410,6 +412,7 @@ impl WorkerGroup { let index_build_job_pool = Arc::new(LocalScheduler::new(config.max_background_index_builds)); let series_index_store = series_index_store_from_config(&config, data_home).await?; + let series_index_disk_usage = Arc::new(AtomicU64::new(0)); 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)); @@ -464,6 +467,7 @@ impl WorkerGroup { write_buffer_manager: write_buffer_manager.clone(), index_build_job_pool: index_build_job_pool.clone(), series_index_store: series_index_store.clone(), + series_index_disk_usage: series_index_disk_usage.clone(), flush_job_pool: flush_job_pool.clone(), compact_job_pool: compact_job_pool.clone(), purge_scheduler: purge_scheduler.clone(), @@ -572,6 +576,7 @@ struct WorkerStarter { compact_job_pool: SchedulerRef, index_build_job_pool: SchedulerRef, series_index_store: Option, + series_index_disk_usage: Arc, flush_job_pool: SchedulerRef, purge_scheduler: SchedulerRef, listener: WorkerListener, @@ -600,16 +605,21 @@ impl WorkerStarter { let (sender, receiver) = mpsc::channel(self.config.worker_channel_size); let running = Arc::new(AtomicBool::new(true)); - let series_index_task_state = self - .series_index_store - .as_ref() - .map(|_| Arc::new(SeriesIndexTaskState::new())); + let (series_index_task_state, series_index_receiver) = if self.series_index_store.is_some() + { + let (state, receiver) = + SeriesIndexTaskState::new(self.id, self.config.worker_channel_size); + (Some(Arc::new(state)), Some(receiver)) + } else { + (None, None) + }; let mut series_index_purger = None; let series_index_handle = self .series_index_store .clone() .zip(series_index_task_state.clone()) - .map(|(store, state)| { + .zip(series_index_receiver) + .map(|((store, state), receiver)| { let (purger, purge_receiver) = series_index_channel(store.clone()); series_index_purger = Some(purger.clone()); spawn_series_index_tasks( @@ -617,9 +627,12 @@ impl WorkerStarter { store, regions.clone(), state, + receiver, self.config.experimental_series_index_bucket_width, purger, purge_receiver, + self.series_index_disk_usage.clone(), + self.config.experimental_series_index_max_size.as_bytes(), self.config.experimental_series_index_maintenance_interval, self.time_provider.clone(), self.config.experimental_enable_range_index, @@ -951,7 +964,7 @@ struct RegionWorkerLoop { index_build_scheduler: IndexBuildScheduler, /// Controls the worker-owned series-index task. series_index_task_state: Option>, - /// Store for companion range indexes deleted by the region SST purger. + /// Local store for series and range indexes managed by reconciliation. series_index_store: Option, series_index_purger: Option, /// Schedules background flush requests. diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index 5a0406c655f..ede17511ad7 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -36,16 +36,12 @@ impl RegionWorkerLoop { info!("Try to close region {}, worker: {}", region_id, self.id); - // If the close request asks for a flush, or the region skips WAL, - // and has data in memtable and region is flushable (like, not in follower state), - // we should flush it before closing to ensure durability. - if (request.flush_on_close || region.skip_wal()) - && !region - .version_control - .current() - .version - .memtables - .is_empty() + // Memtables containing skip-WAL writes must be flushed before closing. + // Followers cannot flush; explicit flush-on-close and region WAL options + // retain their existing behavior. + let version = region.version(); + if (request.flush_on_close || region.skip_wal() || version.memtables.has_skip_wal_writes()) + && !version.memtables.is_empty() && region.is_flushable() { info!("Region {} has pending data, waiting for flush", region_id); @@ -69,7 +65,7 @@ impl RegionWorkerLoop { return; } - // WAL configured or memtable is empty, flush is not necessary. + // No close-time flush is required, or the region cannot flush. self.remove_region(region_id).await; info!("Region {} closed, worker: {}", region_id, self.id); sender.send(Ok(0)) diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index bfdfea5d1d9..1d4442d8d03 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -243,6 +243,9 @@ impl RegionWorkerLoop { flush_semaphore: self.flush_semaphore.clone(), is_staging: region.is_staging(), partition_expr: region.maybe_staging_partition_expr_str(), + durability_barrier: self + .wal + .durability_barrier(region.region_id, ®ion.provider), } } } diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index aa8475a49f1..8aac67dc96f 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -32,7 +32,7 @@ use crate::cache::file_cache::{FileType, IndexKey}; use crate::config::IndexBuildMode; use crate::error::{EditRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, Result}; use crate::manifest::action::{ - RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionTruncate, + RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionTruncate, TruncateKind, }; use crate::memtable::MemtableBuilderProvider; use crate::metrics::WRITE_CACHE_INFLIGHT_DOWNLOAD; @@ -257,13 +257,10 @@ impl RegionWorkerLoop { } else { None }; - let new_mutable = Arc::new( - region - .version() - .memtables - .mutable - .new_with_part_duration(version.compaction_time_window, memtable_builder), - ); + let new_mutable = version + .memtables + .mutable + .new_with_part_duration(version.compaction_time_window, memtable_builder); // Here it assumes the leader has backfilled the partition_expr of the metadata. let metadata = manifest.metadata.clone(); @@ -484,17 +481,31 @@ impl RegionWorkerLoop { let request_sender = self.sender.clone(); let manifest_ctx = region.manifest_ctx.clone(); let is_staging = region.is_staging(); + let durability_barrier = self + .wal + .durability_barrier(region.region_id, ®ion.provider); // Updates manifest in background. common_runtime::spawn_global(async move { + // The truncated entry id becomes the replay frontier, so the WAL + // must be durable through it before the manifest names it. + let durable = match &truncate.kind { + TruncateKind::All { + truncated_entry_id, .. + } => durability_barrier.wait(*truncated_entry_id).await, + _ => Ok(()), + }; // Write region truncated to manifest. let action_list = RegionMetaActionList::with_action(RegionMetaAction::Truncate(truncate.clone())); - let result = manifest_ctx - .update_manifest(RegionLeaderState::Truncating, action_list, is_staging) - .await - .map(|_| ()); + let result = match durable { + Ok(()) => manifest_ctx + .update_manifest(RegionLeaderState::Truncating, action_list, is_staging) + .await + .map(|_| ()), + Err(e) => Err(e), + }; // Sends the result back to the request sender. let truncate_result = TruncateResult { @@ -531,10 +542,12 @@ impl RegionWorkerLoop { let region_id = region.region_id; let request_sender = self.sender.clone(); let manifest_ctx = region.manifest_ctx.clone(); + let durability_barrier = self.wal.durability_barrier(region_id, ®ion.provider); common_runtime::spawn_global(async move { // The frontier moves to the last written entry and sequence, so replaying the - // WAL after a restart skips everything the memtables held. + // WAL after a restart skips everything the memtables held. The WAL must be + // durable through that entry before the manifest names it. let edit = RegionEdit { files_to_add: Vec::new(), files_to_remove: Vec::new(), @@ -545,10 +558,13 @@ impl RegionWorkerLoop { committed_sequence: None, }; let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit)); - let result = manifest_ctx - .update_manifest(RegionLeaderState::Truncating, action_list, false) - .await - .map(|_| ()); + let result = match durability_barrier.wait(discarded_entry_id).await { + Ok(()) => manifest_ctx + .update_manifest(RegionLeaderState::Truncating, action_list, false) + .await + .map(|_| ()), + Err(e) => Err(e), + }; let result = DiscardUnflushedResult { region_id, diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index 3eb45b1ab21..59f42ef4232 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -17,13 +17,14 @@ use std::collections::HashMap; use std::sync::Arc; +use api::v1::region::build_index_request; use common_telemetry::{debug, error, warn}; use store_api::region_request::RegionBuildIndexRequest; use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc; use crate::cache::CacheStrategy; -use crate::error::Result; +use crate::error::{InvalidRequestSnafu, Result, WorkerStoppedSnafu}; use crate::manifest::action::RegionEdit; use crate::metrics::INDEX_PUBLICATION_STALE_TOTAL; use crate::region::version::VersionRef; @@ -68,8 +69,6 @@ impl RegionWorkerLoop { inverted_index_config: self.config.inverted_index.clone(), fulltext_index_config: self.config.fulltext_index.clone(), bloom_filter_index_config: self.config.bloom_filter_index.clone(), - #[cfg(feature = "vector_index")] - vector_index_config: self.config.vector_index.clone(), index_options: version.options.index_options.clone(), intermediate_manager, puffin_manager, @@ -98,18 +97,56 @@ impl RegionWorkerLoop { pub(crate) async fn handle_build_index_request( &mut self, region_id: RegionId, - _req: RegionBuildIndexRequest, + req: RegionBuildIndexRequest, sender: OptionOutputTx, ) { - self.handle_rebuild_index( - BuildIndexRequest { - region_id, - build_type: IndexBuildType::Manual, - file_metas: Vec::new(), - }, - sender, - ) - .await; + match req.options { + Some(build_index_request::Options::SeriesIndex(_)) => { + self.handle_build_series_index_request(region_id, sender); + } + None | Some(build_index_request::Options::SstIndex(_)) => { + self.handle_rebuild_index( + BuildIndexRequest { + region_id, + build_type: IndexBuildType::Manual, + file_metas: Vec::new(), + }, + sender, + ) + .await; + } + } + } + + /// Submits manual reconciliation without blocking the region worker on index I/O. + fn handle_build_series_index_request(&self, region_id: RegionId, mut sender: OptionOutputTx) { + let Some(region) = self.regions.writable_region_or(region_id, &mut sender) else { + return; + }; + let Some(state) = &self.series_index_task_state else { + sender.send( + InvalidRequestSnafu { + region_id, + reason: "series index is disabled; enable experimental_enable_series_index", + } + .fail(), + ); + return; + }; + let receiver = match state.try_reconcile(region) { + Ok(receiver) => receiver, + Err(error) => { + sender.send(Err(error)); + return; + } + }; + let worker_id = self.id; + common_runtime::spawn_global(async move { + let result = receiver + .await + .unwrap_or_else(|_| WorkerStoppedSnafu { id: worker_id }.fail()); + sender.send(result.map(|_| 0)); + }); } pub(crate) async fn handle_rebuild_index( diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 323d73c60ec..41af569c649 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -33,7 +33,7 @@ use crate::metrics::{ WRITE_REJECT_TOTAL, WRITE_ROWS_TOTAL, WRITE_STAGE_ELAPSED, WRITE_STALL_TOTAL, }; use crate::region::{RegionLeaderState, RegionRoleState}; -use crate::region_write_ctx::RegionWriteCtx; +use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest}; use crate::wal::Wal; use crate::worker::RegionWorkerLoop; @@ -126,9 +126,7 @@ impl RegionWorkerLoop { if region_ctxs.len() == 1 { // fast path for single region. let mut region_ctx = region_ctxs.into_values().next().unwrap(); - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; } else { @@ -137,9 +135,7 @@ impl RegionWorkerLoop { .map(|mut region_ctx| { // use tokio runtime to schedule tasks. common_runtime::spawn_global(async move { - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; (region_ctx.put_num, region_ctx.delete_num) }) }) @@ -348,6 +344,7 @@ impl RegionWorkerLoop { ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), + WriteSource::Request, ); e.insert(region_ctx); @@ -492,6 +489,7 @@ impl RegionWorkerLoop { ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), + WriteSource::Request, ); e.insert(region_ctx); @@ -897,6 +895,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); ctx.push_mutation( @@ -946,6 +945,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(1), None, skip_wal)); @@ -959,9 +959,7 @@ mod tests { 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(); + ctx.write_memtables().await; assert_eq!(version_control.committed_sequence(), 1); assert_eq!(version_control.current().last_entry_id, 0); } @@ -995,9 +993,7 @@ mod tests { assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); for region_ctx in region_ctxs.values_mut() { - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; } assert_eq!( @@ -1034,6 +1030,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None, skip_wal)); @@ -1050,9 +1047,7 @@ mod tests { let write_handle = tokio::spawn(async move { let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; }); tokio::time::timeout( diff --git a/src/object-store/Cargo.toml b/src/object-store/Cargo.toml index 412f3563298..15e2069e3bd 100644 --- a/src/object-store/Cargo.toml +++ b/src/object-store/Cargo.toml @@ -8,25 +8,22 @@ license.workspace = true workspace = true [features] +hdfs-object-store = ["dep:hdfs-native", "opendal/services-hdfs-native", "uuid"] mysql-object-store = ["opendal/services-mysql"] services-memory = ["opendal/services-memory"] -testing = ["derive_builder"] +testing = ["derive_builder", "uuid"] [dependencies] -async-trait.workspace = true bytes.workspace = true cap-std.workspace = true -chrono.workspace = true common-base.workspace = true common-error.workspace = true common-macro.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -datafusion_object_store.workspace = true derive_builder = { workspace = true, optional = true } -futures.workspace = true +hdfs-native = { workspace = true, optional = true } humantime-serde.workspace = true -lazy_static.workspace = true opendal = { version = "0.59.2", features = [ "layers-tracing", "layers-prometheus", @@ -43,14 +40,16 @@ reqwest.workspace = true serde.workspace = true snafu.workspace = true tokio.workspace = true -uuid.workspace = true +uuid = { workspace = true, optional = true } [dev-dependencies] anyhow = "1.0" common-telemetry.workspace = true common-test-util.workspace = true +datafusion_object_store.workspace = true +futures.workspace = true +object-store = { workspace = true, features = ["testing"] } object_store_opendal.workspace = true -rand.workspace = true tempfile.workspace = true tokio.workspace = true toml.workspace = true diff --git a/src/object-store/src/config.rs b/src/object-store/src/config.rs index 1308db92fab..2917bdbe817 100644 --- a/src/object-store/src/config.rs +++ b/src/object-store/src/config.rs @@ -12,10 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "hdfs-object-store")] +use std::collections::HashMap; use std::time::Duration; use common_base::readable_size::ReadableSize; use common_base::secrets::{ExposeSecret, SecretString}; +#[cfg(feature = "hdfs-object-store")] +use opendal::services::HdfsNative; #[cfg(feature = "mysql-object-store")] use opendal::services::Mysql; use opendal::services::{Azblob, Gcs, Oss, S3}; @@ -34,6 +38,8 @@ pub enum ObjectStoreConfig { Oss(OssConfig), Azblob(AzblobConfig), Gcs(GcsConfig), + #[cfg(feature = "hdfs-object-store")] + Hdfs(HdfsConfig), #[cfg(feature = "mysql-object-store")] Mysql(MysqlConfig), } @@ -53,6 +59,8 @@ impl ObjectStoreConfig { Self::Oss(_) => "Oss", Self::Azblob(_) => "Azblob", Self::Gcs(_) => "Gcs", + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(_) => "Hdfs", #[cfg(feature = "mysql-object-store")] Self::Mysql(_) => "Mysql", } @@ -72,6 +80,8 @@ impl ObjectStoreConfig { Self::Oss(oss) => &oss.name, Self::Azblob(az) => &az.name, Self::Gcs(gcs) => &gcs.name, + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(hdfs) => &hdfs.name, #[cfg(feature = "mysql-object-store")] Self::Mysql(mysql) => &mysql.name, }; @@ -91,6 +101,8 @@ impl ObjectStoreConfig { Self::Oss(oss) => Some(&oss.cache), Self::Azblob(az) => Some(&az.cache), Self::Gcs(gcs) => Some(&gcs.cache), + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(hdfs) => Some(&hdfs.cache), #[cfg(feature = "mysql-object-store")] Self::Mysql(mysql) => Some(&mysql.cache), } @@ -104,6 +116,8 @@ impl ObjectStoreConfig { Self::Oss(oss) => Some(&mut oss.cache), Self::Azblob(az) => Some(&mut az.cache), Self::Gcs(gcs) => Some(&mut gcs.cache), + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(hdfs) => Some(&mut hdfs.cache), #[cfg(feature = "mysql-object-store")] Self::Mysql(mysql) => Some(&mut mysql.cache), } @@ -289,6 +303,42 @@ impl From<&GcsConnection> for Gcs { } } +/// Connection options for a Hadoop Distributed File System backend. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +#[serde(default)] +#[cfg(feature = "hdfs-object-store")] +pub struct HdfsConnection { + /// Working directory for all object-store operations. + pub root: String, + /// HDFS NameNode URI, for example `hdfs://127.0.0.1:9000`. + pub name_node: String, + /// Additional options passed to the native HDFS client. + pub options: HashMap, +} + +#[cfg(feature = "hdfs-object-store")] +impl From<&HdfsConnection> for HdfsNative { + fn from(connection: &HdfsConnection) -> Self { + let root = util::normalize_dir(&connection.root); + HdfsNative::default() + .root(&root) + .name_node(&connection.name_node) + .options(connection.options.clone()) + } +} + +/// Hadoop Distributed File System object storage configuration. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +#[serde(default)] +#[cfg(feature = "hdfs-object-store")] +pub struct HdfsConfig { + pub name: String, + #[serde(flatten)] + pub connection: HdfsConnection, + #[serde(flatten)] + pub cache: ObjectStorageCacheConfig, +} + #[cfg(feature = "mysql-object-store")] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] #[serde(default)] @@ -411,6 +461,20 @@ mod tests { assert_eq!("test", s3_config.config_name()); assert_eq!("S3", s3_config.provider_name()); + #[cfg(feature = "hdfs-object-store")] + { + let hdfs_config = ObjectStoreConfig::Hdfs(HdfsConfig::default()); + assert_eq!("Hdfs", hdfs_config.config_name()); + assert_eq!("Hdfs", hdfs_config.provider_name()); + + let hdfs_config = ObjectStoreConfig::Hdfs(HdfsConfig { + name: "test".to_string(), + ..Default::default() + }); + assert_eq!("test", hdfs_config.config_name()); + assert_eq!("Hdfs", hdfs_config.provider_name()); + } + #[cfg(feature = "mysql-object-store")] { let mysql_config = ObjectStoreConfig::Mysql(MysqlConfig::default()); @@ -438,6 +502,11 @@ mod tests { assert!(gcs_config.is_object_storage()); let azblob_config = ObjectStoreConfig::Azblob(AzblobConfig::default()); assert!(azblob_config.is_object_storage()); + #[cfg(feature = "hdfs-object-store")] + { + let hdfs_config = ObjectStoreConfig::Hdfs(HdfsConfig::default()); + assert!(hdfs_config.is_object_storage()); + } #[cfg(feature = "mysql-object-store")] { let mysql_config = ObjectStoreConfig::Mysql(MysqlConfig::default()); @@ -445,6 +514,46 @@ mod tests { } } + #[cfg(feature = "hdfs-object-store")] + #[test] + fn test_hdfs_config_serde() { + let config: ObjectStoreConfig = toml::from_str( + r#" +type = "Hdfs" +name = "hdfs-store" +root = "/greptimedb" +name_node = "hdfs://127.0.0.1:9000" + +[options] +"dfs.client.block.write.replace-datanode-on-failure.enable" = "true" +"#, + ) + .unwrap(); + + let ObjectStoreConfig::Hdfs(hdfs_config) = config else { + unreachable!() + }; + + assert_eq!("hdfs-store", hdfs_config.name); + assert_eq!("/greptimedb", hdfs_config.connection.root); + assert_eq!("hdfs://127.0.0.1:9000", hdfs_config.connection.name_node); + assert_eq!( + Some(&"true".to_string()), + hdfs_config + .connection + .options + .get("dfs.client.block.write.replace-datanode-on-failure.enable") + ); + + let serialized = toml::to_string(&hdfs_config).unwrap(); + assert!(serialized.contains("name_node = \"hdfs://127.0.0.1:9000\"")); + assert!( + serialized.contains( + "\"dfs.client.block.write.replace-datanode-on-failure.enable\" = \"true\"" + ) + ); + } + #[cfg(feature = "mysql-object-store")] #[test] fn test_mysql_config_connection_string_serde() { diff --git a/src/object-store/src/factory.rs b/src/object-store/src/factory.rs index 9bdd656a07c..356d3aecb21 100644 --- a/src/object-store/src/factory.rs +++ b/src/object-store/src/factory.rs @@ -15,15 +15,21 @@ use std::{fs, path}; use common_telemetry::info; +#[cfg(feature = "hdfs-object-store")] +use opendal::services::HdfsNative; #[cfg(feature = "mysql-object-store")] use opendal::services::Mysql; use opendal::services::{Fs, Gcs, Oss, S3}; use snafu::prelude::*; +#[cfg(feature = "hdfs-object-store")] +use crate::config::HdfsConfig; #[cfg(feature = "mysql-object-store")] use crate::config::MysqlConfig; use crate::config::{AzblobConfig, FileConfig, GcsConfig, ObjectStoreConfig, OssConfig, S3Config}; use crate::error::{self, Result}; +#[cfg(feature = "hdfs-object-store")] +use crate::layers::HdfsCompatibilityLayer; use crate::services::Azblob; use crate::util::{build_http_context, clean_temp_dir, join_dir, normalize_dir}; use crate::{ATOMIC_WRITE_DIR, OLD_ATOMIC_WRITE_DIR, ObjectStore, util}; @@ -39,11 +45,36 @@ pub async fn new_raw_object_store( ObjectStoreConfig::Oss(oss_config) => new_oss_object_store(oss_config).await, ObjectStoreConfig::Azblob(azblob_config) => new_azblob_object_store(azblob_config).await, ObjectStoreConfig::Gcs(gcs_config) => new_gcs_object_store(gcs_config).await, + #[cfg(feature = "hdfs-object-store")] + ObjectStoreConfig::Hdfs(hdfs_config) => new_hdfs_object_store(hdfs_config).await, #[cfg(feature = "mysql-object-store")] ObjectStoreConfig::Mysql(mysql_config) => new_mysql_object_store(mysql_config).await, } } +/// Creates an object store backed by a native HDFS client. +#[cfg(feature = "hdfs-object-store")] +pub async fn new_hdfs_object_store(hdfs_config: &HdfsConfig) -> Result { + let root = util::normalize_dir(&hdfs_config.connection.root); + info!( + "The HDFS NameNode is: {}, root is: {}", + hdfs_config.connection.name_node, root + ); + + let builder = HdfsNative::from(&hdfs_config.connection); + let compatibility_layer = HdfsCompatibilityLayer::new( + &hdfs_config.connection.name_node, + &hdfs_config.connection.root, + &hdfs_config.connection.options, + ) + .context(error::InitBackendSnafu)?; + let operator = ObjectStore::new(builder) + .context(error::InitBackendSnafu)? + .layer(compatibility_layer); + + Ok(operator) +} + #[cfg(feature = "mysql-object-store")] pub async fn new_mysql_object_store(mysql_config: &MysqlConfig) -> Result { let root = util::normalize_dir(&mysql_config.root); @@ -144,3 +175,33 @@ pub async fn new_s3_object_store(s3_config: &S3Config) -> Result { Ok(operator) } + +#[cfg(all(test, feature = "hdfs-object-store"))] +mod tests { + use opendal::services::HDFS_NATIVE_SCHEME; + + use super::*; + use crate::config::HdfsConnection; + + #[tokio::test] + async fn test_new_hdfs_object_store() { + let config = HdfsConfig { + connection: HdfsConnection { + root: "/greptimedb".to_string(), + name_node: "hdfs://127.0.0.1:9000".to_string(), + ..Default::default() + }, + ..Default::default() + }; + + let store = new_hdfs_object_store(&config).await.unwrap(); + assert_eq!(HDFS_NATIVE_SCHEME, store.info().scheme()); + assert_eq!("/greptimedb/", store.info().root()); + } + + #[tokio::test] + async fn test_new_hdfs_object_store_requires_name_node() { + let result = new_hdfs_object_store(&HdfsConfig::default()).await; + assert!(result.is_err()); + } +} diff --git a/src/object-store/src/layers.rs b/src/object-store/src/layers.rs index cc9fa9f4df5..77c894dbfcd 100644 --- a/src/object-store/src/layers.rs +++ b/src/object-store/src/layers.rs @@ -12,9 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "hdfs-object-store")] +mod hdfs; #[cfg(feature = "testing")] pub mod mock; +#[cfg(feature = "hdfs-object-store")] +pub use hdfs::HdfsCompatibilityLayer; pub use opendal::layers::*; pub use prometheus::build_prometheus_metrics_layer; diff --git a/src/object-store/src/layers/hdfs.rs b/src/object-store/src/layers/hdfs.rs new file mode 100644 index 00000000000..4e00098fbf7 --- /dev/null +++ b/src/object-store/src/layers/hdfs.rs @@ -0,0 +1,551 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt::{self, Debug}; +use std::sync::Arc; + +use hdfs_native::{Client, ClientBuilder}; +use opendal::raw::oio::{Delete as _, Read as _, ReadStream as _, Write as _}; +use opendal::raw::{ + Layer, OpCompose, OpCopy, OpCreateDir, OpDelete, OpList, OpPresign, OpRead, OpRename, OpStat, + OpWrite, RpCreateDir, RpPresign, RpRename, RpStat, Service, ServiceInfo, Servicer, oio, +}; +use opendal::{Buffer, Capability, ErrorKind, Metadata, OperationContext, Result}; +use uuid::Uuid; + +/// Adds atomic writes and streaming copies to the native HDFS backend. +#[derive(Debug, Clone)] +pub struct HdfsCompatibilityLayer { + renamer: AtomicRenamer, +} + +impl HdfsCompatibilityLayer { + /// Creates a compatibility layer for the HDFS connection. + pub fn new( + name_node: &str, + root: &str, + options: &std::collections::HashMap, + ) -> Result { + let mut config = std::collections::HashMap::new(); + let namenodes = name_node + .split(',') + .filter_map(|value| { + let value = value + .trim() + .trim_start_matches("hdfs://") + .trim_end_matches('/'); + (!value.is_empty()).then_some(value) + }) + .collect::>(); + for (index, namenode) in namenodes.iter().enumerate() { + config.insert( + format!("dfs.namenode.rpc-address.nameservice.nn{index}"), + (*namenode).to_string(), + ); + } + config.insert( + "dfs.ha.namenodes.nameservice".to_string(), + (0..namenodes.len()) + .map(|index| format!("nn{index}")) + .collect::>() + .join(","), + ); + config.extend(options.clone()); + let client = ClientBuilder::new() + .with_url("hdfs://nameservice") + .with_config(config) + .build() + .map_err(hdfs_error)?; + Ok(Self { + renamer: AtomicRenamer::Native { + client, + root: opendal::raw::normalize_root(root), + }, + }) + } + + /// Creates a compatibility layer backed by the inner service's rename. + #[cfg(any(test, feature = "testing"))] + pub fn new_for_test() -> Self { + Self { + renamer: AtomicRenamer::Raw, + } + } +} + +#[derive(Debug, Clone)] +enum AtomicRenamer { + Native { + client: Client, + root: String, + }, + #[cfg(any(test, feature = "testing"))] + Raw, +} + +impl AtomicRenamer { + async fn rename( + &self, + _inner: &Servicer, + _ctx: &OperationContext, + from: &str, + to: &str, + ) -> Result<()> { + match self { + Self::Native { client, root } => { + // OpenDAL's HDFS rename removes an existing destination before + // renaming. Use HDFS Rename2 with overwrite to keep replacement atomic. + client + .rename( + &opendal::raw::build_rooted_abs_path(root, from), + &opendal::raw::build_rooted_abs_path(root, to), + true, + ) + .await + .map_err(hdfs_error) + } + #[cfg(any(test, feature = "testing"))] + Self::Raw => _inner + .rename(_ctx, from, to, OpRename::new()) + .await + .map(|_| ()), + } + } +} + +fn hdfs_error(error: hdfs_native::HdfsError) -> opendal::Error { + opendal::Error::new(ErrorKind::Unexpected, "native HDFS operation failed").set_source(error) +} + +impl Layer for HdfsCompatibilityLayer { + fn apply_service(&self, inner: Servicer) -> Servicer { + Arc::new(HdfsCompatibilityService { + inner, + renamer: self.renamer.clone(), + }) + } +} + +struct HdfsCompatibilityService { + inner: Servicer, + renamer: AtomicRenamer, +} + +impl Debug for HdfsCompatibilityService { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HdfsCompatibilityService") + .field("inner", &self.inner) + .finish() + } +} + +/// A writer that publishes non-append writes with an atomic rename. +struct HdfsWriter(HdfsWriterInner); + +enum HdfsWriterInner { + Direct(oio::Writer), + Atomic { + inner: Servicer, + context: OperationContext, + renamer: AtomicRenamer, + writer: Option, + temporary_path: String, + target_path: String, + }, +} + +impl oio::Write for HdfsWriter { + async fn write(&mut self, buffer: Buffer) -> Result<()> { + match &mut self.0 { + HdfsWriterInner::Direct(writer) => writer.write(buffer).await, + HdfsWriterInner::Atomic { writer, .. } => { + writer + .as_mut() + .ok_or_else(writer_unavailable)? + .write(buffer) + .await + } + } + } + + async fn close(&mut self) -> Result { + match &mut self.0 { + HdfsWriterInner::Direct(writer) => writer.close().await, + HdfsWriterInner::Atomic { + inner, + context, + renamer, + writer, + temporary_path, + target_path, + } => { + let mut writer = writer.take().ok_or_else(writer_unavailable)?; + let metadata = match writer.close().await { + Ok(metadata) => metadata, + Err(error) => { + drop(writer); + let _ = delete_path(inner, context, temporary_path).await; + return Err(error); + } + }; + drop(writer); + + if let Err(error) = renamer + .rename(inner, context, temporary_path, target_path) + .await + { + let _ = delete_path(inner, context, temporary_path).await; + return Err(error); + } + + Ok(metadata) + } + } + } + + async fn abort(&mut self) -> Result<()> { + match &mut self.0 { + HdfsWriterInner::Direct(writer) => writer.abort().await, + HdfsWriterInner::Atomic { + inner, + context, + writer, + temporary_path, + .. + } => { + let abort_result = if let Some(mut writer) = writer.take() { + let result = writer.abort().await; + drop(writer); + result + } else { + Ok(()) + }; + let cleanup_result = delete_path(inner, context, temporary_path).await; + + match abort_result { + Err(error) if error.kind() != ErrorKind::Unsupported => Err(error), + _ => cleanup_result, + } + } + } + } +} + +fn writer_unavailable() -> opendal::Error { + opendal::Error::new( + ErrorKind::Unexpected, + "HDFS writer is unavailable after close or abort", + ) +} + +impl Service for HdfsCompatibilityService { + type Reader = oio::Reader; + type Writer = HdfsWriter; + type Lister = oio::Lister; + type Deleter = oio::Deleter; + type Copier = oio::OneShotCopier; + type Composer = oio::Composer; + + fn info(&self) -> ServiceInfo { + self.inner.info() + } + + fn capability(&self) -> Capability { + let mut capability = self.inner.capability(); + capability.copy = true; + capability + } + + async fn create_dir( + &self, + ctx: &OperationContext, + path: &str, + args: OpCreateDir, + ) -> Result { + self.inner.create_dir(ctx, path, args).await + } + + async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result { + self.inner.stat(ctx, path, args).await + } + + fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result { + self.inner.read(ctx, path, args) + } + + fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result { + if args.append() { + return self + .inner + .write(ctx, path, args) + .map(|writer| HdfsWriter(HdfsWriterInner::Direct(writer))); + } + + let temporary_path = temporary_path(path); + let writer = self.inner.write(ctx, &temporary_path, args)?; + Ok(HdfsWriter(HdfsWriterInner::Atomic { + inner: Arc::clone(&self.inner), + context: ctx.clone(), + renamer: self.renamer.clone(), + writer: Some(writer), + temporary_path, + target_path: path.to_string(), + })) + } + + fn copy( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpCopy, + ) -> Result { + if args.if_not_exists() || args.if_match().is_some() { + return Err(opendal::Error::new( + ErrorKind::Unsupported, + "conditional copy is not supported by the HDFS fallback", + )); + } + + let inner = Arc::clone(&self.inner); + let context = ctx.clone(); + let renamer = self.renamer.clone(); + let from = from.to_string(); + let to = to.to_string(); + Ok(oio::OneShotCopier::new(async move { + copy_via_read_write(inner, &context, renamer, &from, &to).await + })) + } + + fn delete(&self, ctx: &OperationContext) -> Result { + self.inner.delete(ctx) + } + + fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result { + self.inner.list(ctx, path, args) + } + + fn compose(&self, ctx: &OperationContext, to: &str, args: OpCompose) -> Result { + self.inner.compose(ctx, to, args) + } + + async fn rename( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpRename, + ) -> Result { + self.inner.rename(ctx, from, to, args).await + } + + async fn presign( + &self, + ctx: &OperationContext, + path: &str, + args: OpPresign, + ) -> Result { + self.inner.presign(ctx, path, args).await + } +} + +async fn copy_via_read_write( + inner: Servicer, + context: &OperationContext, + renamer: AtomicRenamer, + source_path: &str, + target_path: &str, +) -> Result { + let reader = inner.read(context, source_path, OpRead::new())?; + let (_, mut reader) = reader.open(opendal::BytesRange::from(..)).await?; + let temporary_path = temporary_path(target_path); + let mut writer = inner.write(context, &temporary_path, OpWrite::new())?; + + loop { + let buffer = match reader.read().await { + Ok(buffer) => buffer, + Err(error) => { + abort_and_delete(&inner, context, writer, &temporary_path).await; + return Err(error); + } + }; + if buffer.is_empty() { + break; + } + if let Err(error) = writer.write(buffer).await { + abort_and_delete(&inner, context, writer, &temporary_path).await; + return Err(error); + } + } + + let metadata = match writer.close().await { + Ok(metadata) => metadata, + Err(error) => { + drop(writer); + let _ = delete_path(&inner, context, &temporary_path).await; + return Err(error); + } + }; + drop(writer); + + if let Err(error) = renamer + .rename(&inner, context, &temporary_path, target_path) + .await + { + let _ = delete_path(&inner, context, &temporary_path).await; + return Err(error); + } + + Ok(metadata) +} + +async fn abort_and_delete( + inner: &Servicer, + context: &OperationContext, + mut writer: oio::Writer, + path: &str, +) { + let _ = writer.abort().await; + drop(writer); + let _ = delete_path(inner, context, path).await; +} + +async fn delete_path(inner: &Servicer, context: &OperationContext, path: &str) -> Result<()> { + let mut deleter = inner.delete(context)?; + deleter.delete(path, OpDelete::new()).await?; + deleter.close().await +} + +// TODO(fengjiachun): Clean up temporary files left behind after a process crash. +fn temporary_path(path: &str) -> String { + let suffix = format!(".greptime-{}.tmp", Uuid::new_v4()); + match path.rsplit_once('/') { + Some((parent, name)) => format!("{parent}/.{name}{suffix}"), + None => format!(".{path}{suffix}"), + } +} + +#[cfg(test)] +mod tests { + use opendal::services::Fs; + use opendal::{Operator, Writer}; + use tempfile::TempDir; + + use super::*; + + fn test_store() -> (TempDir, Operator) { + let directory = tempfile::tempdir().unwrap(); + let store = Operator::new(Fs::default().root(directory.path().to_str().unwrap())) + .unwrap() + .layer(HdfsCompatibilityLayer::new_for_test()); + (directory, store) + } + + #[tokio::test] + async fn test_service_operations() { + let (_directory, store) = test_store(); + store.create_dir("data/").await.unwrap(); + store.write("data/source", "contents").await.unwrap(); + assert_eq!(8, store.stat("data/source").await.unwrap().content_length()); + assert_eq!( + b"onte", + store + .read_with("data/source") + .range(1..5) + .await + .unwrap() + .to_bytes() + .as_ref() + ); + store.rename("data/source", "data/target").await.unwrap(); + let entries = store.list("data/").await.unwrap(); + assert!(entries.iter().any(|entry| entry.path() == "data/target")); + assert!(!store.exists("data/source").await.unwrap()); + store.delete("data/target").await.unwrap(); + assert!(!store.exists("data/target").await.unwrap()); + } + + #[tokio::test] + async fn test_atomic_write_keeps_old_data_after_abort() { + let (_directory, store) = test_store(); + store.write("manifest.json", "old").await.unwrap(); + + let mut writer: Writer = store.writer("manifest.json").await.unwrap(); + writer.write("new").await.unwrap(); + writer.abort().await.unwrap(); + + assert_eq!( + b"old", + store + .read("manifest.json") + .await + .unwrap() + .to_bytes() + .as_ref() + ); + assert!( + store + .list("") + .await + .unwrap() + .iter() + .all(|entry| !entry.path().contains(".greptime-")) + ); + } + + #[tokio::test] + async fn test_atomic_write_replaces_on_close() { + let (_directory, store) = test_store(); + store.write("manifest.json", "old").await.unwrap(); + store.write("manifest.json", "new").await.unwrap(); + + assert_eq!( + b"new", + store + .read("manifest.json") + .await + .unwrap() + .to_bytes() + .as_ref() + ); + assert!( + store + .list("") + .await + .unwrap() + .iter() + .all(|entry| !entry.path().contains(".greptime-")) + ); + } + + #[tokio::test] + async fn test_copy_fallback_streams_to_target() { + let (_directory, store) = test_store(); + store.write("source.parquet", "contents").await.unwrap(); + store + .copy("source.parquet", "target.parquet") + .await + .unwrap(); + + assert_eq!( + b"contents", + store + .read("target.parquet") + .await + .unwrap() + .to_bytes() + .as_ref() + ); + } +} diff --git a/src/object-store/src/lib.rs b/src/object-store/src/lib.rs index 05c344485c6..6c8956445d8 100644 --- a/src/object-store/src/lib.rs +++ b/src/object-store/src/lib.rs @@ -24,11 +24,13 @@ pub mod error; pub mod factory; pub mod layers; pub mod manager; -mod metrics; pub mod secure_fs; +#[cfg(feature = "testing")] pub mod test_util; pub mod util; +#[cfg(feature = "hdfs-object-store")] +pub use config::HdfsConnection; pub use config::{AzblobConnection, GcsConnection, OssConnection, S3Connection}; /// The default object cache directory name. diff --git a/src/object-store/src/metrics.rs b/src/object-store/src/metrics.rs deleted file mode 100644 index 28aee460f6c..00000000000 --- a/src/object-store/src/metrics.rs +++ /dev/null @@ -1,15 +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. - -//! object-store metrics diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 18afa3ac820..feaaa42fdf2 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -274,6 +274,7 @@ impl Service for SecureFsBackend { path, args, file: None, + synced: false, }) } @@ -387,6 +388,29 @@ struct SecureFsWriter { path: PathBuf, args: OpWrite, file: Option, + synced: bool, +} + +#[derive(Debug)] +struct UnsyncedOverwrite { + flush_error: Option, +} + +impl fmt::Display for UnsyncedOverwrite { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "overwrite file was not synced") + } +} + +impl std::error::Error for UnsyncedOverwrite { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.flush_error.as_ref().map(|error| error as _) + } +} + +/// Whether an abort error proves an opened overwrite has not been synced. +pub fn is_unsynced_overwrite_abort(error: &Error) -> bool { + std::error::Error::source(error).is_some_and(|source| source.is::()) } impl SecureFsWriter { @@ -438,10 +462,18 @@ impl oio::Write for SecureFsWriter { } async fn close(&mut self) -> Result { - let file = self.ensure_file().await?; - file.flush().await.map_err(new_std_io_error)?; - file.sync_all().await.map_err(new_std_io_error)?; - let metadata = file.metadata().await.map_err(new_std_io_error)?; + { + let file = self.ensure_file().await?; + file.flush().await.map_err(new_std_io_error)?; + file.sync_all().await.map_err(new_std_io_error)?; + } + self.synced = true; + let metadata = self + .ensure_file() + .await? + .metadata() + .await + .map_err(new_std_io_error)?; let mut builder = MetadataBuilder::file(metadata.len()); builder.last_modified(Timestamp::try_from( metadata.modified().map_err(new_std_io_error)?, @@ -450,10 +482,43 @@ impl oio::Write for SecureFsWriter { } async fn abort(&mut self) -> Result<()> { - Err(Error::new( + // Tokio writes may finish in the blocking pool after write_all returns. + let flush = match self.file.as_mut() { + Some(file) => file.flush().await.map_err(new_std_io_error), + None => Ok(()), + }; + if self.args.if_not_exists() { + // A failed exclusive create owns no file. Once data is synced, preserve + // potentially committed output for the caller's deliberate retry. + if let Some(file) = self.file.take() { + drop(file); + if !self.synced { + let root = self.root.clone(); + let path = self.path.clone(); + let cleanup = + common_runtime::spawn_blocking_global(move || root.dir.remove_file(path)) + .await + .map_err(new_task_join_error) + .and_then(|result| result.map_err(new_std_io_error)); + return flush.and(cleanup); + } + } + return flush; + } + if self.file.is_none() { + return flush; + } + let error = Error::new( ErrorKind::Unsupported, "filesystem writes cannot be aborted without atomic writes", - )) + ); + if !self.synced { + return Err(error.set_source(UnsyncedOverwrite { + flush_error: flush.err(), + })); + } + flush?; + Err(error) } } @@ -800,16 +865,181 @@ mod tests { } #[tokio::test] - async fn test_writer_abort_is_unsupported_without_atomic_write() { - let temp_dir = create_temp_dir("secure_fs_writer_abort"); + async fn test_conditional_abort_only_removes_owned_partial_file() { + let temp_dir = create_temp_dir("secure_fs_conditional_abort"); let operator = SecureFsRoot::open(temp_dir.path()) .unwrap() .build_operator(); - let mut writer = operator.writer("partial").await.unwrap(); - writer.write(Bytes::from_static(b"partial")).await.unwrap(); + for started in [false, true] { + let mut writer = operator + .writer_with("partial") + .if_not_exists(true) + .await + .unwrap(); + if started { + writer.write(Bytes::from_static(b"partial")).await.unwrap(); + } + writer.abort().await.unwrap(); + assert!(!operator.exists("partial").await.unwrap()); + } + } - let error = writer.abort().await.unwrap_err(); + #[tokio::test] + async fn test_overwrite_abort_preserves_unopened_destination() { + use std::path::PathBuf; - assert_eq!(ErrorKind::Unsupported, error.kind()); + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + + let temp_dir = create_temp_dir("secure_fs_unopened_abort"); + std::fs::write(temp_dir.path().join("existing"), b"original").unwrap(); + let mut writer = super::SecureFsWriter { + root: SecureFsRoot::open(temp_dir.path()).unwrap(), + path: PathBuf::from("existing"), + args: OpWrite::default(), + file: None, + synced: false, + }; + + writer.abort().await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join("existing")).unwrap(), + b"original" + ); + } + + #[cfg(target_os = "linux")] + #[tokio::test] + async fn test_abort_after_background_write_error() { + use std::path::PathBuf; + + use opendal::options::WriteOptions; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + use tokio::io::AsyncWriteExt; + + for if_not_exists in [true, false] { + let temp_dir = create_temp_dir("secure_fs_flush_error_abort"); + let root = SecureFsRoot::open(temp_dir.path()).unwrap(); + let (args, _) = OpWrite::from_options( + &root.build_operator().info().capability(), + WriteOptions { + if_not_exists, + ..Default::default() + }, + ) + .unwrap(); + let mut writer = super::SecureFsWriter { + root, + path: PathBuf::from("partial"), + args, + file: None, + synced: false, + }; + writer.ensure_file().await.unwrap(); + let mut failing_file = tokio::fs::OpenOptions::new() + .write(true) + .open("/dev/full") + .await + .unwrap(); + failing_file.write_all(b"partial").await.unwrap(); + writer.file = Some(failing_file); + + let error = writer.abort().await.unwrap_err(); + if if_not_exists { + assert!(!temp_dir.path().join("partial").exists()); + } else { + assert_eq!(error.kind(), opendal::ErrorKind::Unsupported); + assert!(std::error::Error::source(&error).is_some()); + } + } + } + + #[cfg(target_os = "linux")] + #[tokio::test] + async fn test_conditional_abort_after_close_flush_error() { + use std::path::PathBuf; + + use opendal::options::WriteOptions; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + use tokio::io::AsyncWriteExt; + + let temp_dir = create_temp_dir("secure_fs_close_flush_error"); + let root = SecureFsRoot::open(temp_dir.path()).unwrap(); + let (args, _) = OpWrite::from_options( + &root.build_operator().info().capability(), + WriteOptions { + if_not_exists: true, + ..Default::default() + }, + ) + .unwrap(); + let mut writer = super::SecureFsWriter { + root, + path: PathBuf::from("partial"), + args, + file: None, + synced: false, + }; + writer.ensure_file().await.unwrap(); + let mut failing_file = tokio::fs::OpenOptions::new() + .write(true) + .open("/dev/full") + .await + .unwrap(); + failing_file.write_all(b"partial").await.unwrap(); + writer.file = Some(failing_file); + + assert!(writer.close().await.is_err()); + writer.abort().await.unwrap(); + assert!(!temp_dir.path().join("partial").exists()); + } + + #[test] + fn test_writer_abort_drains_before_reporting_unsupported() { + use std::path::PathBuf; + + use opendal::Buffer; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + + use super::SecureFsWriter; + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap() + .block_on(async { + let temp_dir = create_temp_dir("secure_fs_writer_abort"); + let mut writer = SecureFsWriter { + root: SecureFsRoot::open(temp_dir.path()).unwrap(), + path: PathBuf::from("partial"), + args: OpWrite::default(), + file: None, + synced: false, + }; + writer.ensure_file().await.unwrap(); + let (started, ready) = tokio::sync::oneshot::channel(); + let (release, blocked) = std::sync::mpsc::channel(); + let blocking = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + blocked.recv().unwrap(); + }); + ready.await.unwrap(); + writer.write(Buffer::from("partial")).await.unwrap(); + let abort = writer.abort(); + tokio::pin!(abort); + let pending = futures::poll!(&mut abort).is_pending(); + release.send(()).unwrap(); + assert!(pending); + assert_eq!(ErrorKind::Unsupported, abort.await.unwrap_err().kind()); + blocking.await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join("partial")).unwrap(), + b"partial" + ); + }); } } diff --git a/src/object-store/src/test_util.rs b/src/object-store/src/test_util.rs index af68978b0e4..936ecba5bac 100644 --- a/src/object-store/src/test_util.rs +++ b/src/object-store/src/test_util.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::env; - use crate::{ObjectStore, Result}; /// Temp folder for object store test @@ -35,30 +33,3 @@ impl TempFolder { self.store.delete_with(&self.path).recursive(true).await } } - -/// Test s3 config from environment variables -#[derive(Debug)] -pub struct TestS3Config { - pub root: String, - pub access_key_id: String, - pub secret_access_key: String, - pub bucket: String, - pub region: Option, -} - -/// Returns s3 test config, return None if not found. -pub fn s3_test_config() -> Option { - if let Ok(b) = env::var("GT_S3_BUCKET") - && !b.is_empty() - { - return Some(TestS3Config { - root: uuid::Uuid::new_v4().to_string(), - access_key_id: env::var("GT_S3_ACCESS_KEY_ID").ok()?, - secret_access_key: env::var("GT_S3_ACCESS_KEY").ok()?, - bucket: env::var("GT_S3_BUCKET").ok()?, - region: Some(env::var("GT_S3_REGION").ok()?), - }); - } - - None -} diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index 9cf03fe42c5..dc2482d1722 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -284,12 +284,23 @@ async fn test_s3_backend() -> Result<()> { let root = uuid::Uuid::new_v4().to_string(); - let builder = S3::default() + let mut builder = S3::default() .root(&root) .access_key_id(&env::var("GT_S3_ACCESS_KEY_ID")?) .secret_access_key(&env::var("GT_S3_ACCESS_KEY")?) .region(&env::var("GT_S3_REGION")?) .bucket(&bucket); + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") { + builder = builder.endpoint(&endpoint); + } + + // Honors an S3-compatible endpoint (MinIO in CI) so the test does not + // fall back to resolving the bucket against real AWS. + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") + && !endpoint.is_empty() + { + builder = builder.endpoint(&endpoint); + } let store = ObjectStore::new(builder).unwrap(); let store = object_store::util::with_instrument_layers(store, false); @@ -298,6 +309,7 @@ async fn test_s3_backend() -> Result<()> { test_object_crud(&store).await?; test_object_list(&store).await?; test_object_list_start_after(&store).await?; + test_conditional_creation(&store).await?; assert_opendal_metrics(); guard.remove_all().await?; } @@ -305,6 +317,74 @@ async fn test_s3_backend() -> Result<()> { Ok(()) } +async fn test_conditional_creation(store: &ObjectStore) -> Result<()> { + const PART: usize = 8 * 1024 * 1024; + for size in [32, 2 * PART + 17] { + let payload = Bytes::from(vec![42; size]); + for preexisting in [false, true] { + let path = format!("conditional-{size}-{preexisting}"); + if preexisting { + store.write(&path, "original").await?; + } + let attempt = |value: Bytes| { + let path = &path; + async move { + let mut writer = store + .writer_with(path) + .if_not_exists(true) + .concurrent(1) + .chunk(PART) + .await + .unwrap(); + let result = async { + for offset in (0..value.len()).step_by(PART) { + writer + .write(value.slice(offset..(offset + PART).min(value.len()))) + .await?; + } + writer.close().await + } + .await; + if result.is_err() { + writer.abort().await.unwrap(); + } + result + } + }; + let (a, b) = tokio::join!(attempt(payload.clone()), attempt(payload.clone())); + assert_eq!( + usize::from(a.is_ok()) + usize::from(b.is_ok()), + usize::from(!preexisting) + ); + for error in [a.err(), b.err()].into_iter().flatten() { + assert!( + matches!( + error.kind(), + object_store::ErrorKind::ConditionNotMatch + | object_store::ErrorKind::AlreadyExists + ), + "{error:?}" + ); + } + let expected = if preexisting { + Bytes::from_static(b"original") + } else { + payload.clone() + }; + assert_eq!(store.read(&path).await?.to_bytes(), expected); + store.delete(&path).await?; + } + } + Ok(()) +} + +#[tokio::test] +async fn test_secure_fs_conditional_creation() -> Result<()> { + let dir = TempDir::new()?; + let store = object_store::secure_fs::SecureFsRoot::open(dir.path())?.build_operator(); + test_conditional_creation(&store).await +} + #[tokio::test] async fn test_oss_backend() -> Result<()> { common_telemetry::init_default_ut_logging(); diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 62853a18f61..49cbd617567 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -46,3 +46,5 @@ the default `meter-macros/noop` build and active metering when changing this pat Finite requests split internally use `admit_write` / `admit_row_insert_batches` before dispatch. Their database-scoped `QueryContext` admission prevents a second row debit while `do_request` still records WCU for the actual writes. +Async table-batcher outputs carry storage completions in `OutputMeta`; keep +them when combining submissions so callers can distinguish admission from success. diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index ea1bc990f74..72f953e0f89 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -22,7 +22,6 @@ workspace = true ahash.workspace = true api.workspace = true arrow.workspace = true -arrow-ipc.workspace = true async-stream.workspace = true async-trait.workspace = true bytes.workspace = true @@ -66,7 +65,6 @@ object-store.workspace = true partition.workspace = true percent-encoding.workspace = true prometheus.workspace = true -prost.workspace = true prost-types = { workspace = true, optional = true } query.workspace = true regex.workspace = true diff --git a/src/operator/src/batcher.rs b/src/operator/src/batcher.rs index 4c5e9892c0b..bc9ef86b34e 100644 --- a/src/operator/src/batcher.rs +++ b/src/operator/src/batcher.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use arrow::record_batch::RecordBatch; use async_trait::async_trait; +use common_query::Output; use session::context::QueryContextRef; use table::metadata::TableInfoRef; use tokio::sync::OwnedSemaphorePermit; @@ -29,13 +30,15 @@ pub trait PendingRowsBatcher: Send + Sync { /// Acquires one slot per original request, shared by all of its table submissions. async fn acquire(&self) -> Result>; - /// Waits for the submitted rows to be written, retaining the slot through completion. + /// Submits rows according to the acknowledgement policy, retaining the slot until + /// writing completes even when the response acknowledges queue admission only. /// Cancelling the response wait does not retract an already enqueued submission. + /// Early acknowledgements retain actual storage results in the output metadata. async fn submit( &self, table_info: TableInfoRef, batch: RecordBatch, ctx: QueryContextRef, permit: Arc, - ) -> Result; + ) -> Result; } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index f1f57a64479..f8d58f5d173 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -140,6 +140,24 @@ impl Inserter { return Ok(0); } + // Bulk storage paths may discard unknown columns instead of rejecting them. + for field in record_batch.schema_ref().fields() { + ensure!( + table_info + .meta + .schema + .column_schema_by_name(field.name()) + .is_some(), + error::InvalidInsertRequestSnafu { + reason: format!( + "Column '{}' not found in table '{}'", + field.name(), + table_info.full_table_name() + ), + } + ); + } + // The zero value is WCU, not bytes. Bulk writes have no WCU accounting; // preserve that behavior while admitting their rows before dispatch. write_meter!(MeterRecord::new( diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index b28af78d383..d69277a1c90 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -24,9 +24,10 @@ use api::v1::column_def::{options_from_column_schema, try_as_column_schema}; use api::v1::{ AddColumn, AddColumns, AlterDatabaseExpr, AlterTableExpr, Analyzer, ColumnDataType, ColumnDataTypeExtension, CreateFlowExpr, CreateTableExpr, CreateViewExpr, DropColumn, - DropColumns, DropDefaults, ExpireAfter, FulltextBackend as PbFulltextBackend, ModifyColumnType, + DropColumns, DropDefaults, ExpireAfter, FulltextBackend as PbFulltextBackend, + JsonSettings as PbJsonSettings, JsonTypeHint as PbJsonTypeHint, ModifyColumnType, ModifyColumnTypes, RenameTable, SemanticType, SetDatabaseOptions, SetDefaults, SetFulltext, - SetIndex, SetIndexes, SetInverted, SetSkipping, SetTableOptions, + SetIndex, SetIndexes, SetInverted, SetJsonSettings, SetSkipping, SetTableOptions, SkippingIndexType as PbSkippingIndexType, TableName, UnsetDatabaseOptions, UnsetFulltext, UnsetIndex, UnsetIndexes, UnsetInverted, UnsetSkipping, UnsetTableOptions, set_index, unset_index, @@ -36,6 +37,7 @@ use common_error::ext::BoxedError; use common_grpc_expr::util::ColumnExpr; use common_time::Timezone; use datafusion::sql::planner::object_name_to_table_reference; +use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, @@ -66,8 +68,8 @@ use sql::statements::create::{ CreateView, TableConstraint, }; use sql::statements::{ - OptionMap, column_to_schema, concrete_data_type_to_sql_data_type, - sql_column_def_to_grpc_column_def, sql_data_type_to_concrete_data_type, value_to_sql_value, + OptionMap, column_to_schema, concrete_data_type_to_sql_data_type, sql_col_to_grpc_col_def, + sql_data_type_to_concrete_data_type, value_to_sql_value, }; use sql::util::extract_tables_from_query; use store_api::mito_engine_options::{COMPACTION_OVERRIDE, COMPACTION_TYPE}; @@ -782,6 +784,29 @@ pub(crate) fn to_repartition_request( }) } +fn json_settings_to_proto(settings: JsonSettings) -> Result { + let (type_hints, max_auto_expanded_paths) = settings.into_parts(); + let type_hints = type_hints + .into_iter() + .map(|hint| { + let (data_type, datatype_extension) = ColumnDataTypeWrapper::try_from(hint.data_type) + .map(|w| w.to_parts()) + .context(ColumnDataTypeSnafu)?; + + Ok(PbJsonTypeHint { + path: hint.path, + data_type: data_type as i32, + datatype_extension, + }) + }) + .collect::>>()?; + + Ok(PbJsonSettings { + type_hints, + max_auto_expanded_paths, + }) +} + /// Converts a SQL alter table statement into a gRPC alter table expression. pub(crate) fn to_alter_table_expr( alter_table: AlterTable, @@ -803,12 +828,10 @@ pub(crate) fn to_alter_table_expr( add_columns: add_columns .into_iter() .map(|add_column| { - let column_def = sql_column_def_to_grpc_column_def( - &add_column.column_def, - Some(&query_ctx.timezone()), - ) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; + let column_def = + sql_col_to_grpc_col_def(&add_column.column, Some(&query_ctx.timezone())) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; if is_interval_type(&column_def.data_type()) { return NotSupportedSnafu { feat: "Add column with interval type", @@ -830,12 +853,15 @@ pub(crate) fn to_alter_table_expr( } => { let target_type = sql_data_type_to_concrete_data_type(&target_type).context(ParseSqlSnafu)?; + + // Currently disallow modify column type to json2. 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)?; @@ -853,6 +879,20 @@ pub(crate) fn to_alter_table_expr( }], }) } + AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + } => { + let settings = match json2_options { + Some(options) => options.build_json_settings().context(ParseSqlSnafu)?, + None => datatypes::json::JsonSettings::new_v2(), + }; + + AlterTableKind::SetJsonSettings(SetJsonSettings { + column_name: column_name.value, + settings: Some(json_settings_to_proto(settings)?), + }) + } AlterTableOperation::DropColumn { name } => AlterTableKind::DropColumns(DropColumns { drop_columns: vec![DropColumn { name: name.value.clone(), @@ -1150,6 +1190,7 @@ mod tests { use std::collections::HashMap; use api::v1::{SetDatabaseOptions, UnsetDatabaseOptions}; + use datatypes::extension::json::is_json2_extension_type; use datatypes::value::Value; use session::context::{QueryContext, QueryContextBuilder}; use sql::dialect::GreptimeDbDialect; @@ -1159,6 +1200,55 @@ mod tests { use super::*; + #[test] + fn test_to_alter_add_json2_column_preserves_extension() { + let sql = "ALTER TABLE monitor ADD COLUMN payload JSON2;"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let AlterTableKind::AddColumns(AddColumns { add_columns, .. }) = expr.kind.unwrap() else { + unreachable!() + }; + let column_def = add_columns[0].column_def.as_ref().unwrap(); + let column_schema = try_as_column_schema(column_def).unwrap(); + let schema = Schema::new(vec![column_schema]); + + assert!(is_json2_extension_type(&schema.arrow_schema().fields()[0])); + } + + #[test] + fn test_to_alter_add_json2_column_preserves_settings() { + let sql = "ALTER TABLE monitor ADD COLUMN payload JSON2(max_auto_expanded_paths = 0, service STRING, nested.value BIGINT)"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let AlterTableKind::AddColumns(AddColumns { add_columns, .. }) = expr.kind.unwrap() else { + unreachable!() + }; + let column_schema = + try_as_column_schema(add_columns[0].column_def.as_ref().unwrap()).unwrap(); + let metadata: datatypes::extension::json::JsonMetadata = + serde_json::from_str(&column_schema.metadata()["ARROW:extension:metadata"]).unwrap(); + let settings = metadata.json_settings(); + assert_eq!(Some(0), settings.max_auto_expanded_paths()); + assert_eq!(2, settings.type_hints().len()); + assert_eq!(vec!["service"], settings.type_hints()[0].path); + assert_eq!(vec!["nested", "value"], settings.type_hints()[1].path); + } + #[test] fn test_create_flow_tql_expr() { let sql = r#" @@ -1752,31 +1842,48 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; } #[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(); + fn test_to_alter_set_json_settings_expr() { + let sql = "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(); + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let kind = expr.kind.unwrap(); - assert!(matches!(err, crate::error::Error::NotSupported { .. })); - assert_eq!( - "Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type", - err.to_string() - ); - } + let AlterTableKind::SetJsonSettings(modify) = kind else { + unreachable!() + }; + assert_eq!("payload", modify.column_name); + let settings = modify.settings.as_ref().unwrap(); + assert_eq!(Some(100), settings.max_auto_expanded_paths); + assert_eq!(1, settings.type_hints.len()); + assert_eq!(["service"], &settings.type_hints[0].path[..]); + + let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2;"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let kind = expr.kind.unwrap(); + + let AlterTableKind::SetJsonSettings(modify) = kind else { + unreachable!() + }; + let settings = modify.settings.as_ref().unwrap(); + assert_eq!(Some(100), settings.max_auto_expanded_paths); + assert!(settings.type_hints.is_empty()); } #[test] diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index afc903621e2..c7f60c0883c 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -13,19 +13,23 @@ // limitations under the License. use std::collections::BTreeMap; -use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::v1::alter_table_expr::Kind; -use api::v1::column_def::options_from_skipping; +use api::v1::column_def::{options_from_skipping, try_as_column_def}; use api::v1::region::{ InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests, RegionRequestHeader, }; +use api::v1::value::ValueData; use api::v1::{ AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests, RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; +use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ @@ -46,6 +50,8 @@ use common_query::native_histogram::{is_native_histogram_value_type, native_hist use common_query::prelude::{greptime_timestamp, greptime_value}; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; +use common_time::Timestamp; +use common_time::timestamp::TimeUnit; use datatypes::schema::SkippingIndexOptions; use futures_util::future; use meter_core::data::MeterRecord; @@ -100,6 +106,10 @@ pub struct Inserter { /// per-request `auto_create_table` hint. When `true`, the hint still applies. auto_create_table: bool, pending_rows_batcher: Option>, + /// Rows handed to detached flow mirror tasks that have not finished yet. + /// Bounds how much row data in-flight mirror writes may keep alive, see + /// [`MAX_MIRROR_PENDING_ROWS`]. + mirror_pending_rows: Arc, } pub type InserterRef = Arc; @@ -156,6 +166,132 @@ pub struct InstantAndNormalInsertRequests { } impl Inserter { + /// Checks the assumptions of the logical bulk path without changing tables. + /// Unsupported requests retain ordinary insertion, including schema policy, + /// defaults, instant TTL and row-based Flow delivery. + pub async fn can_batch_metric_rows( + &self, + requests: &RowInsertRequests, + ctx: &QueryContextRef, + physical_table: &str, + ) -> Result { + if self.auto_create_disabled_reason(ctx)?.is_some() || ctx.extension(TTL_KEY).is_some() { + return Ok(false); + } + for request in &requests.inserts { + // The logical bulk encoder only supports scalar metric schemas; + // any time index unit is accepted — requests are converted to the + // destination table's unit during batch alignment. + // Check new tables too, before catalog lookup or schema changes. + if request.rows.as_ref().is_some_and(|rows| { + rows.schema.iter().any(|column| { + column.datatype_extension.is_some() + || !matches!( + ColumnDataType::try_from(column.datatype), + Ok(ColumnDataType::TimestampSecond + | ColumnDataType::TimestampMillisecond + | ColumnDataType::TimestampMicrosecond + | ColumnDataType::TimestampNanosecond + | ColumnDataType::Float64 + | ColumnDataType::String) + ) + }) + }) { + return Ok(false); + } + + let Some(table) = self + .get_table( + ctx.current_catalog(), + &ctx.current_schema(), + &request.table_name, + ) + .await? + else { + continue; + }; + let info = table.table_info(); + if info.meta.engine != METRIC_ENGINE_NAME + || info.is_ttl_instant_table() + || info + .meta + .options + .extra_options + .get(LOGICAL_TABLE_METADATA_KEY) + .map(String::as_str) + != Some(physical_table) + || info + .meta + .schema + .column_schemas() + .iter() + .any(|column| column.default_constraint().is_some()) + { + return Ok(false); + } + // Physical metric tags are nullable even when their logical schema + // is not. Arrow alignment is stricter than ordinary metric insertion. + if info + .meta + .primary_key_indices + .iter() + .any(|&index| !info.meta.schema.column_schemas()[index].is_nullable()) + { + return Ok(false); + } + // The current Flow cache does not distinguish streaming and batch + // flows. Keep all Flow sources on the row-based delivery path. + match self.table_flownode_set_cache.get(info.table_id()).await { + Ok(None) => {} + Ok(Some(flows)) if flows.is_empty() => {} + _ => return Ok(false), + } + } + Ok(true) + } + + /// Meters an original logical-table request before bulk routing, without + /// cloning its rows or changing the request boundary used for accounting. + pub async fn meter_row_inserts( + requests: &mut RowInsertRequests, + ctx: &QueryContextRef, + ) -> Result { + let metered = InstantAndNormalInsertRequests { + normal_requests: RegionInsertRequests { + requests: requests + .inserts + .iter_mut() + .map(|request| RegionInsertRequest { + rows: request.rows.take(), + ..Default::default() + }) + .collect(), + }, + instant_requests: RegionInsertRequests::default(), + }; + let cost = write_meter!( + ctx.current_catalog(), + ctx.current_schema(), + metered, + ctx.write_rows_to_admit( + ctx.current_catalog(), + &ctx.current_schema(), + count_insert_rows(&metered)? + ), + ctx.channel() as u8 + ) + .await + .context(WriteRejectedSnafu); + for (request, region) in requests + .inserts + .iter_mut() + .zip(metered.normal_requests.requests) + { + request.rows = region.rows; + } + cost + } + pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, @@ -170,6 +306,7 @@ impl Inserter { table_flownode_set_cache, auto_create_table, pending_rows_batcher: None, + mirror_pending_rows: Arc::new(AtomicU64::new(0)), } } @@ -308,6 +445,7 @@ impl Inserter { statement_executor, accommodate_existing_schema, is_single_value, + None, ) .await?; @@ -417,11 +555,14 @@ impl Inserter { // Observe every table completion even when another table fails. future::join_all(submissions).await }; - let affected_rows = results.into_iter().sum::>()?; - Ok(Output::new( - OutputData::AffectedRows(affected_rows), - OutputMeta::new_with_cost(write_cost as _), - )) + let mut affected_rows = 0; + let mut meta = OutputMeta::new_with_cost(write_cost as _); + for result in results { + let output = result?; + affected_rows += output.extract_rows_and_cost().0; + meta.write_completions.extend(output.meta.write_completions); + } + Ok(Output::new(OutputData::AffectedRows(affected_rows), meta)) } /// Handles row inserts request with metric engine. @@ -444,10 +585,21 @@ impl Inserter { validate_column_count_match(&requests)?; // check and create physical table - self.create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor) + let physical_table_ref = self + .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor) .await?; - // check and create logical tables + // check and create logical tables; `create_or_alter_tables_on_demand` + // aligns each request's time index unit with the unit of the table it + // targets, inside its existing table lookups: existing tables keep + // their own unit (which matches the physical table they are bound + // to), and new tables use the selected physical table's unit (from + // `physical_table_ref`). Ingestion endpoints encode timestamps in a + // fixed unit (prometheus remote write always uses millisecond; OTLP + // keeps nanosecond precision on the metric engine path), and the + // metric engine requires each logical table's requests to match its + // time index unit. Narrowing conversions truncate the sub-unit part + // (floor), following `Timestamp::convert_to`. let CreateAlterTableResult { instant_table_ids, table_infos, @@ -459,6 +611,7 @@ impl Inserter { statement_executor, true, true, + table_time_index_unit(&physical_table_ref), ) .await?; let name_to_info = table_infos @@ -661,7 +814,17 @@ impl Inserter { .chain(instant_requests.requests.iter()), ) .await?; - flow_mirror_task.detach(self.node_manager.clone())?; + let has_instant_rows = instant_requests.requests.iter().any(|request| { + request + .rows + .as_ref() + .is_some_and(|rows| !rows.rows.is_empty()) + }); + flow_mirror_task.detach( + self.node_manager.clone(), + self.mirror_pending_rows.clone(), + has_instant_rows, + )?; // Write requests to datanode and wait for response let write_tasks = self @@ -764,6 +927,84 @@ impl Inserter { ) } + /// Adds missing columns from a bulk stream's schema and returns the refreshed table. + /// Call once when initializing the stream, before writing its first batch. + /// Does not infer new nested or dictionary columns. + pub async fn ensure_bulk_insert_schema( + &self, + table: TableRef, + request_schema: &ArrowSchema, + ctx: &QueryContextRef, + statement_executor: &StatementExecutor, + ) -> Result { + let table_info = table.table_info(); + if self.auto_create_disabled_reason(ctx)?.is_some() + && !Self::is_auto_create_exempt_private_table(&table_info.schema_name, &table_info.name) + { + return Ok(table); + } + + let table_schema = table.schema(); + let schema = request_schema + .fields() + .iter() + .filter(|field| table_schema.column_schema_by_name(field.name()).is_none()) + .map(|field| { + let data_type = field.data_type(); + // Dictionary values can reach the same infallible child-type conversion + // as nested types, even when Arrow's is_nested() returns false. + ensure!( + !data_type.is_nested() && !matches!(data_type, ArrowDataType::Dictionary(..)), + crate::error::NotSupportedSnafu { + feat: format!( + "automatically adding bulk insert column '{}' with type {:?}", + field.name(), + data_type + ), + } + ); + let column = datatypes::schema::ColumnSchema::try_from(field.as_ref()) + .context(crate::error::ConvertSchemaSnafu)?; + // Arrow fields do not carry primary-key semantics. New columns are + // fields, unless explicitly marked as a time index. + let column_def = + try_as_column_def(&column, false).context(crate::error::ColumnDataTypeSnafu)?; + Ok(ColumnSchema { + column_name: column_def.name, + datatype: column_def.data_type, + semantic_type: column_def.semantic_type, + datatype_extension: column_def.datatype_extension, + options: column_def.options, + }) + }) + .collect::>>()?; + let mut request = RowInsertRequest { + table_name: table_info.name.clone(), + rows: Some(Rows { + schema, + rows: Vec::new(), + }), + }; + let Some(alter_expr) = + self.get_alter_table_expr_on_demand(&mut request, &table, ctx, false, false, true)? + else { + return Ok(table); + }; + + statement_executor + .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter) + .await?; + self.get_table( + &table_info.catalog_name, + &table_info.schema_name, + &table_info.name, + ) + .await? + .with_context(|| TableNotFoundSnafu { + table_name: table_info.full_table_name(), + }) + } + /// Ensures a trace table has the request-global schema without requiring a /// padded data row to drive on-demand creation or alteration. When /// `alter_existing` is false, a table created after planning is left for the @@ -792,6 +1033,7 @@ impl Inserter { statement_executor, false, false, + None, ) .await?; Ok(()) @@ -809,6 +1051,13 @@ impl Inserter { /// custom schema, and then inserts data with endpoints that have default schema setting, like prometheus /// remote write. This will modify the `RowInsertRequests` in place. /// `is_single_value` indicates whether the default schema only contains single value column so we can accommodate it. + /// + /// `align_time_index_unit` is the selected physical metric table's time + /// index unit; passing `Some` (metric engine path only) rewrites each + /// request's time index column, inside this function's existing table + /// lookups (no extra catalog access): existing destination tables are + /// converted to their own unit, and new tables to the given unit. + #[allow(clippy::too_many_arguments)] async fn create_or_alter_tables_on_demand( &self, requests: &mut RowInsertRequests, @@ -817,6 +1066,7 @@ impl Inserter { statement_executor: &StatementExecutor, accommodate_existing_schema: bool, is_single_value: bool, + align_time_index_unit: Option, ) -> Result { let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND .with_label_values(&[auto_create_table_type.as_str()]) @@ -837,7 +1087,7 @@ impl Inserter { && !has_auto_create_exempt_table { let mut instant_table_ids = HashSet::new(); - for req in &requests.inserts { + for req in &mut requests.inserts { let table = match self.get_table(catalog, &schema, &req.table_name).await? { Some(table) => table, // System-defined table: created canonically by the system, @@ -857,6 +1107,15 @@ impl Inserter { .fail(); } }; + // Metric path: an existing destination table keeps its own + // time index unit (it may be bound to another physical + // table than the one selected by this request). + if align_time_index_unit.is_some() + && let Some(rows) = req.rows.as_mut() + && let Some(target_unit) = table_time_index_unit(&table) + { + convert_rows_time_unit(rows, target_unit)?; + } let table_info = table.table_info(); if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) { validate_trace_table_model(&table_info, ctx)?; @@ -892,6 +1151,15 @@ impl Inserter { if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); } + // Metric path: an existing destination table keeps its + // own time index unit (it may be bound to another + // physical table than the one selected by this request). + if align_time_index_unit.is_some() + && let Some(rows) = req.rows.as_mut() + && let Some(target_unit) = table_time_index_unit(&table) + { + convert_rows_time_unit(rows, target_unit)?; + } if auto_create_allowed && let Some(alter_expr) = self.get_alter_table_expr_on_demand( req, @@ -937,6 +1205,14 @@ impl Inserter { .fail(); } None => { + // Metric path: a new table uses the selected physical + // table's unit; convert before the create expression is + // derived from the request schema. + if let Some(physical_unit) = align_time_index_unit + && let Some(rows) = req.rows.as_mut() + { + convert_rows_time_unit(rows, physical_unit)?; + } let semantic_index = per_table_semantics .get_or_insert_with(|| parse_per_table_semantic_index(ctx)) .as_ref(); @@ -1135,17 +1411,16 @@ impl Inserter { ctx: &QueryContextRef, physical_table: String, statement_executor: &StatementExecutor, - ) -> Result<()> { + ) -> Result { let catalog_name = ctx.current_catalog(); let schema_name = ctx.current_schema(); // check if exist - if self + if let Some(table) = self .get_table(catalog_name, &schema_name, &physical_table) .await? - .is_some() { - return Ok(()); + return Ok(table); } // Gate here too, otherwise a disabled switch would still leak the physical table. @@ -1197,7 +1472,7 @@ impl Inserter { .await; match res { - Ok(_) => Ok(()), + Ok(table) => Ok(table), Err(err) => { error!(err; "Failed to create table {table_reference}"); Err(err) @@ -1471,6 +1746,84 @@ fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool { ) } +/// Returns the table's time index unit, if any. A metric table without a +/// timestamp column is left alone by the unit alignment: the metric engine +/// rejects it anyway. +fn table_time_index_unit(table: &TableRef) -> Option { + table + .table_info() + .meta + .schema + .timestamp_column() + .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit())) +} + +fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> { + let Some(ts_index) = rows + .schema + .iter() + .position(|col| col.semantic_type == SemanticType::Timestamp as i32) + else { + return Ok(()); + }; + let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype) + .ok() + .and_then(api::helper::timestamp_unit) + else { + return Ok(()); + }; + if source_unit == target_unit { + return Ok(()); + } + + rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32; + // Timestamp columns never carry a datatype extension. + rows.schema[ts_index].datatype_extension = None; + + // Note: the schema is rewritten before the rows are converted, so an + // overflow error mid-batch leaves this request half-converted. That is + // harmless: the error aborts the whole insert request. + // + // `validate_column_count_match` guarantees every row carries exactly one + // value per schema column, so the time index position is directly in + // bounds; no per-value search is needed. + for row in &mut rows.rows { + debug_assert_eq!(row.values.len(), rows.schema.len()); + let value = &mut row.values[ts_index]; + let Some(value_data) = value.value_data.take() else { + continue; + }; + value.value_data = + convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?; + } + Ok(()) +} + +fn convert_timestamp_value_data( + value_data: ValueData, + source_unit: TimeUnit, + target_unit: TimeUnit, + column_index: usize, +) -> Result> { + let timestamp = match value_data { + ValueData::TimestampSecondValue(v) => Timestamp::new_second(v), + ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v), + ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v), + ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v), + // Null or non-timestamp value; nothing to convert. + other => return Ok(Some(other)), + }; + let converted = timestamp + .convert_to(target_unit) + .with_context(|| InvalidInsertRequestSnafu { + reason: format!( + "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}", + timestamp.value() + ), + })?; + Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data) +} + fn table_is_native_histogram(table: &TableRef) -> bool { let mut fields = table.field_columns(); let Some(col) = fields.next() else { @@ -1671,9 +2024,65 @@ struct CreateAlterTableResult { table_infos: HashMap>, } +/// Upper bound of rows buffered by detached flow mirror tasks on this frontend. +/// +/// Rows retained by in-flight mirror RPCs count against this bound. If ordinary +/// persisted source-table mirroring exceeds it, that mirror batch is best-effort +/// dropped while its datanode write continues. Instant-TTL batches fail retryably +/// on temporary saturation; oversized batches fail non-retryably so the client +/// can reduce the batch size. +const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000; + +/// How often at most each mirror event is reported. +const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10); +const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX; + +struct MirrorLog { + last_log_millis: AtomicU64, + events: AtomicU64, +} + +impl MirrorLog { + const fn new() -> Self { + Self { + last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED), + events: AtomicU64::new(0), + } + } + + /// Returns the count since the previous report, including this event, to + /// the single caller that claims the interval. + fn claim_report(&self, now_millis: u64) -> Option { + self.events.fetch_add(1, Ordering::Relaxed); + let last = self.last_log_millis.load(Ordering::Relaxed); + if last != MIRROR_LOG_NEVER_REPORTED + && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64 + { + return None; + } + if self + .last_log_millis + .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed) + .is_err() + { + return None; + } + Some(self.events.swap(0, Ordering::Relaxed)) + } +} + +static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new(); +static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new(); + +/// Milliseconds since the first call, so the drop log is rate limited on a +/// monotonic clock instead of the wall clock, which can jump. +fn mirror_drop_log_millis() -> u64 { + static START: LazyLock = LazyLock::new(Instant::now); + START.elapsed().as_millis() as u64 +} + struct FlowMirrorTask { requests: HashMap, - num_rows: usize, } impl FlowMirrorTask { @@ -1683,7 +2092,6 @@ impl FlowMirrorTask { ) -> Result { let mut src_table_reqs: HashMap, RegionInsertRequests)>> = HashMap::new(); - let mut num_rows = 0; for req in requests { let table_id = RegionId::from_u64(req.region_id).table_id(); @@ -1707,11 +2115,6 @@ impl FlowMirrorTask { if !peers.is_empty() { let mut reqs = RegionInsertRequests::default(); reqs.requests.push(req.clone()); - num_rows += reqs - .requests - .iter() - .map(|r| r.rows.as_ref().unwrap().rows.len()) - .sum::(); src_table_reqs.insert(table_id, Some((peers, reqs))); } else { // insert a empty entry to avoid repeat query @@ -1747,16 +2150,69 @@ impl FlowMirrorTask { } } - Ok(Self { - requests: inserts, - num_rows, - }) + Ok(Self { requests: inserts }) } - fn detach(self, node_manager: NodeManagerRef) -> Result<()> { - crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(self.num_rows as i64); + /// Rows of row data this task keeps alive once detached: the payload it + /// spawns per peer, not the source requests it was built from. + /// + /// A source table mapped to several flownodes clones its requests per peer + /// and every clone owns its own copy of the rows, so reserving from + /// `self.requests` is what the per-peer shares released by [`Self::detach`] + /// sum to. Counting the source requests instead under-reserves whenever a + /// batch holds several requests for the same source table. + fn pending_rows(&self) -> u64 { + self.requests.values().map(region_inserts_rows).sum() + } + + fn detach( + self, + node_manager: NodeManagerRef, + mirror_pending_rows: Arc, + has_instant_rows: bool, + ) -> Result<()> { + // Reserve the cloned per-peer payload before spawning. A full budget + // drops normal mirrors or rejects instant writes before datanode dispatch. + let num_rows = self.pending_rows(); + if num_rows == 0 { + return Ok(()); + } + if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 { + return InvalidInsertRequestSnafu { + reason: format!( + "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size", + MAX_MIRROR_PENDING_ROWS + ), + } + .fail(); + } + + let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows); + if pending > MAX_MIRROR_PENDING_ROWS as u64 { + release_mirror_pending_rows(&mirror_pending_rows, num_rows); + if has_instant_rows { + return Err(meter_core::collect::WriteRejected::new( + "flow mirror pending rows limit exceeded for instant-TTL insert", + )) + .context(WriteRejectedSnafu); + } + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows); + if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) { + warn!( + "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report", + pending, MAX_MIRROR_PENDING_ROWS, events + ); + } + return Ok(()); + } + for (peer, inserts) in self.requests { + // Each spawned task releases exactly its own share of the + // reservation. The shares are computed over the same per-peer + // payload as `pending_rows`, so they sum to the reservation. + let peer_rows = region_inserts_rows(&inserts); let node_manager = node_manager.clone(); + let mirror_pending_rows = mirror_pending_rows.clone(); common_runtime::spawn_global(async move { let result = node_manager .flownode(&peer) @@ -1769,12 +2225,18 @@ impl FlowMirrorTask { Ok(resp) => { let affected_rows = resp.affected_rows; crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows); - crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(affected_rows as _); } Err(err) => { - error!(err; "Failed to insert data into flownode {}", peer); + if let Some(events) = + MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis()) + { + error!(err; flownode_id = peer.id, flownode_addr = %peer.addr, "Failed to insert data into flownode ({} total mirror failures across peers since last report)", events); + } } } + // Release the reservation on both the success and the failure + // path, otherwise a broken flownode would fill the budget forever. + release_mirror_pending_rows(&mirror_pending_rows, peer_rows); }); } @@ -1782,17 +2244,62 @@ impl FlowMirrorTask { } } +/// Rows of row data carried by `inserts`. A request without a `rows` payload +/// carries nothing to keep buffered and contributes zero. +fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 { + inserts + .requests + .iter() + .filter_map(|req| req.rows.as_ref()) + .map(|rows| rows.rows.len() as u64) + .sum() +} + +/// Reserves `rows` of the pending mirror budget and returns the resulting +/// pending row count. The caller must [`release_mirror_pending_rows`] exactly +/// these rows once they are no longer buffered, whether it spawns the batch or +/// drops it. +fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 { + let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows; + // Keep the gauge in step by applying the same increment, never by setting an + // absolute value: another task may move the budget in between. + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64); + pending +} + +/// Releases `rows` of the pending mirror budget. +/// +/// Every release mirrors the reservation it hands back, so a well-behaved +/// release never subtracts more than is pending; the saturating update only +/// keeps a duplicated release from wrapping the shared budget around. +fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) { + let prev = mirror_pending_rows + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| { + Some(pending.saturating_sub(rows)) + }) + .unwrap_or_else(|prev| prev); + // Subtract from the gauge with the opposite sign of + // [`reserve_mirror_pending_rows`]; reading the budget here and `set`-ting it + // would publish a stale value whenever another task releases concurrently. + let released = rows.min(prev); + if released > 0 { + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64); + } +} + #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicU64, Ordering}; use api::helper::ColumnDataTypeWrapper; + use api::v1::flow::FlowResponse; use api::v1::helper::{field_column_schema, time_index_column_schema}; - use api::v1::{RowInsertRequest, Rows, Value}; + use api::v1::{Row, RowInsertRequest, Rows, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_meta::cache::new_table_flownode_set_cache; use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler; - use common_meta::test_util::MockDatanodeManager; + use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager}; use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; use common_query::prelude::{greptime_native_histogram, set_default_prefix}; use datatypes::data_type::ConcreteDataType; @@ -1855,6 +2362,206 @@ mod tests { )) } + fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef { + let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![ + ColumnSchema::new( + greptime_timestamp(), + ConcreteDataType::timestamp_datatype(unit), + false, + ) + .with_time_index(true), + ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true), + ]) + .unwrap() + .build() + .unwrap(); + let meta = TableMetaBuilder::empty() + .schema(Arc::new(schema)) + .primary_key_indices(vec![]) + .value_indices(vec![1]) + .engine("metric") + .next_column_id(0) + .options(Default::default()) + .created_on(Default::default()) + .build() + .unwrap(); + let info = Arc::new( + TableInfoBuilder::default() + .table_id(1) + .table_version(0) + .name("greptime_physical_table") + .schema_name(DEFAULT_SCHEMA_NAME) + .catalog_name(DEFAULT_CATALOG_NAME) + .desc(None) + .table_type(TableType::Base) + .meta(meta) + .build() + .unwrap(), + ); + Arc::new(table::Table::new( + info, + table::metadata::FilterPushDownType::Unsupported, + Arc::new(DummyDataSource), + )) + } + + fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest { + ms_row_insert_request_named("my_metric", timestamp_ms) + } + + fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest { + RowInsertRequest { + table_name: table.to_string(), + rows: Some(Rows { + schema: vec![ + time_index_column_schema( + greptime_timestamp(), + ColumnDataType::TimestampMillisecond, + ), + field_column_schema(greptime_value(), ColumnDataType::Float64), + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)), + }, + Value { + value_data: Some(ValueData::F64Value(1.0)), + }, + ], + }], + }), + } + } + + /// Converts the time index column of each request to `target_unit`. + fn convert_row_insert_requests_time_unit( + requests: &mut RowInsertRequests, + target_unit: TimeUnit, + ) -> Result<()> { + for request in &mut requests.inserts { + let Some(rows) = request.rows.as_mut() else { + continue; + }; + convert_rows_time_unit(rows, target_unit)?; + } + Ok(()) + } + + #[test] + fn test_convert_row_insert_requests_time_unit_noop_when_matching() { + let mut requests = RowInsertRequests { + inserts: vec![ms_row_insert_request(123)], + }; + convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap(); + let rows = requests.inserts[0].rows.as_ref().unwrap(); + assert_eq!( + rows.schema[0].datatype, + ColumnDataType::TimestampMillisecond as i32 + ); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMillisecondValue(123)) + )); + } + + #[test] + fn test_convert_row_insert_requests_time_unit_widens_losslessly() { + let mut requests = RowInsertRequests { + inserts: vec![ms_row_insert_request(123)], + }; + convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap(); + let rows = requests.inserts[0].rows.as_ref().unwrap(); + assert_eq!( + rows.schema[0].datatype, + ColumnDataType::TimestampMicrosecond as i32 + ); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMicrosecondValue(123_000)) + )); + } + + #[test] + fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() { + // 123_456_789 ns floors to 123_456 us and 123 ms; negative values + // floor towards negative infinity, matching `Timestamp::convert_to`. + let requests = |value_ns: i64| RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "my_metric".to_string(), + rows: Some(Rows { + schema: vec![ + time_index_column_schema( + greptime_timestamp(), + ColumnDataType::TimestampNanosecond, + ), + field_column_schema(greptime_value(), ColumnDataType::Float64), + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampNanosecondValue(value_ns)), + }, + Value { + value_data: Some(ValueData::F64Value(1.0)), + }, + ], + }], + }), + }], + }; + + let mut reqs = requests(123_456_789); + convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap(); + let rows = reqs.inserts[0].rows.as_ref().unwrap(); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMicrosecondValue(123_456)) + )); + + let mut reqs = requests(123_456_789); + convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap(); + let rows = reqs.inserts[0].rows.as_ref().unwrap(); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMillisecondValue(123)) + )); + + let mut reqs = requests(-123_456_789); + convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap(); + let rows = reqs.inserts[0].rows.as_ref().unwrap(); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMicrosecondValue(-123_457)) + )); + } + + #[test] + fn test_convert_row_insert_requests_time_unit_overflow() { + let mut requests = RowInsertRequests { + inserts: vec![ms_row_insert_request(i64::MAX)], + }; + let err = + convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err(); + assert!(err.to_string().contains("overflows"), "{err}"); + } + + #[test] + fn test_table_time_index_unit() { + assert_eq!( + table_time_index_unit(&make_metric_physical_table_ref_with_time_unit( + TimeUnit::Microsecond + )), + Some(TimeUnit::Microsecond) + ); + assert_eq!( + table_time_index_unit(&make_metric_physical_table_ref_with_time_unit( + TimeUnit::Millisecond + )), + Some(TimeUnit::Millisecond) + ); + } + #[tokio::test] async fn test_accommodate_existing_schema_and_reject_kind_changes() { let ts_name = "my_ts"; @@ -2313,19 +3020,67 @@ mod tests { ); assert_eq!(ctx.channel(), Channel::Prometheus); } - let attempts = meter.attempts.lock().unwrap(); - let totals = attempts + { + let attempts = meter.attempts.lock().unwrap(); + let totals = attempts + .iter() + .map(|r| (r.schema.as_str(), r.rows, r.value)) + .collect::>(); + assert_eq!( + totals, + if enabled { + vec![("a", 4, 0), ("b", 2, 0)] + } else { + vec![] + } + ); + } + meter.attempts.lock().unwrap().clear(); + meter.reject.store(false, Ordering::Relaxed); + let ctx = Arc::new(QueryContext::with_channel( + CATALOG, + "logical", + Channel::Otlp, + )); + let admitted = admit_write(2, &ctx).await.unwrap(); + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "metric".to_string(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + }], + }; + let original = requests.clone(); + let cost = Inserter::meter_row_inserts(&mut requests, &admitted) + .await + .unwrap(); + assert_eq!(cost, if enabled { 17 } else { 0 }); + assert_eq!(requests, original); + let records = meter + .attempts + .lock() + .unwrap() .iter() - .map(|r| (r.schema.as_str(), r.rows, r.value)) + .map(|record| (record.rows, record.value)) .collect::>(); assert_eq!( - totals, + records, if enabled { - vec![("a", 4, 0), ("b", 2, 0)] + vec![(2, 0), (0, 17)] } else { vec![] } ); + meter.reject.store(true, Ordering::Relaxed); + let result = Inserter::meter_row_inserts(&mut requests, &admitted).await; + if enabled { + assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited); + } else { + assert_eq!(result.unwrap(), 0); + } + assert_eq!(requests, original); } #[test] @@ -2550,7 +3305,7 @@ mod tests { _batch: arrow::record_batch::RecordBatch, _ctx: QueryContextRef, _permit: Arc, - ) -> Result { + ) -> Result { panic!("empty writes must not submit a batch") } } @@ -2570,6 +3325,142 @@ mod tests { ) } + #[tokio::test] + async fn test_logical_batcher_eligibility() { + use catalog::RegisterTableRequest; + use catalog::memory::MemoryCatalogManager; + use common_meta::instruction::{CacheIdent, CreateFlow}; + use common_meta::kv_backend::KvBackendRef; + use common_meta::kv_backend::memory::MemoryKvBackend; + use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder}; + let requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "test_table".to_string(), + rows: None, + }], + }; + let original = + make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype()) + .table_info(); + for case in [ + "eligible", + "physical", + "ordinary", + "instant", + "disabled", + "hint", + "flow", + "required_tag", + "default", + ] { + let mut info = (*original).clone(); + info.meta.engine = METRIC_ENGINE_NAME.to_string(); + info.meta.options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + "physical".to_string(), + ); + let mut ctx = QueryContext::arc().fork(); + match case { + "physical" => { + info.meta + .options + .extra_options + .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string()); + } + "ordinary" => info.meta.engine = "mito".to_string(), + "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant), + "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"), + "required_tag" | "default" => { + let mut columns = info.meta.schema.column_schemas().to_vec(); + let mut tag = ColumnSchema::new( + "tag", + ConcreteDataType::string_datatype(), + case != "required_tag", + ); + if case == "default" { + tag = tag + .with_default_constraint(Some(ColumnDefaultConstraint::null_value())) + .unwrap(); + } + columns.push(tag); + info.meta.schema = Arc::new( + SchemaBuilder::try_from_columns(columns) + .unwrap() + .build() + .unwrap(), + ); + info.meta.primary_key_indices = vec![2]; + } + _ => {} + } + let catalog = MemoryCatalogManager::with_default_setup(); + let table = Arc::new(table::Table::new( + Arc::new(info), + table::metadata::FilterPushDownType::Unsupported, + Arc::new(DummyDataSource), + )); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "test_table".to_string(), + table_id: 1, + table, + }) + .unwrap(); + let mut inserter = batcher_test_inserter().await; + inserter.catalog_manager = catalog; + inserter.auto_create_table = case != "disabled"; + let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default()); + inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache( + String::new(), + Cache::new(10), + kv_backend, + )); + if case == "flow" { + inserter + .table_flownode_set_cache + .invalidate(&[CacheIdent::CreateFlow(CreateFlow { + flow_id: 1, + source_table_ids: vec![1], + partition_to_peer_mapping: vec![(0, Peer::empty(1))], + })]) + .await + .unwrap(); + } + assert_eq!( + inserter + .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical") + .await + .unwrap(), + case == "eligible", + "{case}" + ); + } + } + + #[tokio::test] + async fn test_batcher_meter_preserves_request() { + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "sample".to_string(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row { + values: vec![Value { + value_data: Some(api::v1::value::ValueData::F64Value(1.5)), + }], + }], + }), + }], + }; + let expected = requests.clone(); + Inserter::meter_row_inserts(&mut requests, &QueryContext::arc()) + .await + .unwrap(); + assert_eq!(requests, expected); + } + #[tokio::test] async fn test_instant_table_bypasses_batcher() { let batcher: Arc = Arc::new(UnexpectedBatcher); @@ -2623,4 +3514,695 @@ mod tests { .unwrap(); assert!(matches!(output.data, OutputData::AffectedRows(0))); } + + /// Flownode handler that keeps mirror inserts in flight until the test + /// releases the gate, so the pending mirror budget can be observed + /// deterministically. + #[derive(Clone)] + struct GatedFlownodeHandler { + gate: Arc, + affected_rows: u64, + } + + #[async_trait::async_trait] + impl MockFlownodeHandler for GatedFlownodeHandler { + async fn handle_inserts( + &self, + _peer: &Peer, + _requests: api::v1::region::InsertRequests, + ) -> common_meta::error::Result { + let _permit = self.gate.acquire().await.unwrap(); + Ok(FlowResponse { + affected_rows: self.affected_rows, + ..Default::default() + }) + } + } + + fn flownode_peer() -> Peer { + Peer { + id: 1, + addr: "127.0.0.1:4001".to_string(), + } + } + + #[derive(Clone)] + struct RecordingNodeManager { + datanode_dispatch: tokio::sync::mpsc::UnboundedSender, + flownode_dispatch: tokio::sync::mpsc::UnboundedSender, + gates: Arc>>>, + fail_flownode: bool, + } + + #[async_trait::async_trait] + impl common_meta::node_manager::DatanodeManager for RecordingNodeManager { + async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef { + Arc::new(RecordingNode { + peer: peer.clone(), + manager: self.clone(), + }) + } + } + + #[async_trait::async_trait] + impl common_meta::node_manager::FlownodeManager for RecordingNodeManager { + async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef { + Arc::new(RecordingNode { + peer: peer.clone(), + manager: self.clone(), + }) + } + } + + struct RecordingNode { + peer: Peer, + manager: RecordingNodeManager, + } + + #[async_trait::async_trait] + impl common_meta::node_manager::Datanode for RecordingNode { + async fn handle( + &self, + _request: api::v1::region::RegionRequest, + ) -> common_meta::error::Result { + let rows = match _request.body.as_ref() { + Some(api::v1::region::region_request::Body::Inserts(requests)) => requests + .requests + .iter() + .filter_map(|request| request.rows.as_ref()) + .map(|rows| rows.rows.len()) + .sum(), + _ => 0, + }; + let _ = self.manager.datanode_dispatch.send(rows); + Ok(api::region::RegionResponse::new(rows)) + } + + async fn handle_query( + &self, + _request: common_query::request::QueryRequest, + ) -> common_meta::error::Result { + unreachable!() + } + } + + #[async_trait::async_trait] + impl common_meta::node_manager::Flownode for RecordingNode { + async fn handle( + &self, + _request: api::v1::flow::FlowRequest, + ) -> common_meta::error::Result { + unreachable!() + } + + async fn handle_inserts( + &self, + _request: api::v1::region::InsertRequests, + ) -> common_meta::error::Result { + let _ = self.manager.flownode_dispatch.send(self.peer.clone()); + if self.manager.fail_flownode { + return Err(common_meta::error::UnexpectedSnafu { + err_msg: "test flownode failure".to_string(), + } + .build()); + } + let gate = self + .manager + .gates + .lock() + .unwrap() + .get(&self.peer.id) + .cloned(); + if let Some(gate) = gate { + let _permit = gate.acquire().await.unwrap(); + } + Ok(FlowResponse::default()) + } + + async fn handle_mark_window_dirty( + &self, + _request: api::v1::flow::DirtyWindowRequests, + ) -> common_meta::error::Result { + unreachable!() + } + } + + fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap { + HashMap::from_iter([( + peer.clone(), + RegionInsertRequests { + requests: vec![RegionInsertRequest { + region_id: RegionId::new(1, 1).as_u64(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row { values: vec![] }; num_rows], + }), + ..Default::default() + }], + }, + )]) + } + + /// Serializes the mirror metric tests below: the metrics are process-wide + /// singletons, so these tests would otherwise observe each other's updates + /// when the test binary runs tests in parallel (they are isolated per + /// process only under nextest). + static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + + async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> { + MIRROR_METRIC_TEST_LOCK.lock().await + } + + // The tests below share the process-wide mirror metrics, so they take + // `MIRROR_METRIC_TEST_LOCK` and assert on the difference their own budget + // publishes instead of on absolute values. + #[tokio::test] + async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() { + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + + let _guard = lock_mirror_metrics().await; + let kv_backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(kv_backend.clone()).await; + let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel(); + let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel(); + let node_manager = Arc::new(RecordingNodeManager { + datanode_dispatch: datanode_tx, + flownode_dispatch: flownode_tx, + gates: Arc::new(std::sync::Mutex::new(HashMap::new())), + fail_flownode: false, + }); + let flow_cache = Cache::new(10); + let flow_cache_backend = prepare_mocked_backend().await; + let inserter = Inserter::new( + catalog::memory::MemoryCatalogManager::new(), + partition_manager, + node_manager, + Arc::new(new_table_flownode_set_cache( + String::new(), + flow_cache.clone(), + flow_cache_backend, + )), + true, + ); + let peer = flownode_peer(); + inserter + .table_flownode_set_cache + .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow( + common_meta::instruction::CreateFlow { + flow_id: 1, + source_table_ids: vec![1, 2], + partition_to_peer_mapping: vec![(0, peer)], + }, + )]) + .await + .unwrap(); + let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter()); + normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string(); + normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string(); + let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter()); + instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant); + let table_infos = + HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]); + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + let insert = |table_id, num_rows| RegionInsertRequests { + requests: vec![RegionInsertRequest { + region_id: RegionId::new(table_id, 1).as_u64(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); num_rows], + }), + ..Default::default() + }], + }; + let make_request = |normal, instant| InstantAndNormalInsertRequests { + normal_requests: if normal == 0 { + RegionInsertRequests::default() + } else { + insert(1, normal) + }, + instant_requests: if instant == 0 { + RegionInsertRequests::default() + } else { + insert(2, instant) + }, + }; + + let pending = inserter.mirror_pending_rows.clone(); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64); + pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed); + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + for request in [make_request(0, 1), make_request(1, 1)] { + let error = inserter + .do_request(request, &table_infos, &ctx) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RateLimited); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert!(datanode_rx.try_recv().is_err()); + assert!(flownode_rx.try_recv().is_err()); + assert_eq!( + pending.load(Ordering::Relaxed), + MAX_MIRROR_PENDING_ROWS as u64 + ); + } + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + ); + + for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] { + if pending_before == 0 { + pending.store(0, Ordering::Relaxed); + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64); + } + let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + for (normal, instant) in [ + (0, MAX_MIRROR_PENDING_ROWS + 1), + (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1), + ] { + let error = inserter + .do_request(make_request(normal, instant), &table_infos, &ctx) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + assert_eq!(error.retry_hint(), RetryHint::NonRetryable); + let message = error.to_string(); + assert!(message.contains("1000001"), "{message}"); + assert!(message.contains("1000000"), "{message}"); + assert!(message.contains("reduce the batch size"), "{message}"); + assert!(datanode_rx.try_recv().is_err()); + assert!(flownode_rx.try_recv().is_err()); + assert_eq!(pending.load(Ordering::Relaxed), pending_before); + assert_eq!( + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + gauge_before_rejection + ); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + ); + } + if pending_before == 0 { + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64); + pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed); + } + } + + let result = inserter + .do_request(make_request(1, 0), &table_infos, &ctx) + .await + .unwrap(); + assert!(matches!(result.data, OutputData::AffectedRows(1))); + assert_eq!(1, datanode_rx.try_recv().unwrap()); + assert!(flownode_rx.try_recv().is_err()); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + 1 + ); + pending.store(0, Ordering::Relaxed); + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64); + assert_eq!( + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + gauge_before + ); + } + + #[tokio::test] + async fn flow_mirror_new_counts_regions_and_peer_clones() { + use common_meta::instruction::{CacheIdent, CreateFlow}; + + let peer1 = flownode_peer(); + let peer2 = Peer { + id: 2, + addr: "127.0.0.1:4002".to_string(), + }; + let cache_backend = prepare_mocked_backend().await; + let cache = Arc::new(new_table_flownode_set_cache( + String::new(), + Cache::new(10), + cache_backend, + )); + cache + .invalidate(&[CacheIdent::CreateFlow(CreateFlow { + flow_id: 1, + source_table_ids: vec![1], + partition_to_peer_mapping: vec![(0, peer1), (1, peer2)], + })]) + .await + .unwrap(); + let request = |region_id| RegionInsertRequest { + region_id, + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row { values: vec![] }; 3], + }), + ..Default::default() + }; + let requests = [ + request(RegionId::new(1, 0).as_u64()), + request(RegionId::new(1, 1).as_u64()), + ]; + let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap(); + assert_eq!(12, task.pending_rows()); + assert_eq!(2, task.requests.len()); + } + + #[tokio::test] + async fn flow_mirror_dropped_when_pending_exceeds_limit() { + let _guard = lock_mirror_metrics().await; + let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64)); + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + let num_rows = 100; + + let task = FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), num_rows), + }; + // The budget is already full, so nothing may be spawned for this batch. + task.detach( + Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)), + pending.clone(), + false, + ) + .unwrap(); + + assert_eq!( + num_rows as u64, + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before + ); + // The dropped batch must not leak its reservation. + assert_eq!( + MAX_MIRROR_PENDING_ROWS as u64, + pending.load(Ordering::Relaxed) + ); + // The reservation taken before the limit check is released again, so the + // gauge nets out and keeps matching the pending budget. + assert_eq!( + gauge_before, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + "a dropped batch must not change the pending gauge" + ); + + let task = FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), num_rows), + }; + let error = task + .detach( + Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)), + pending.clone(), + true, + ) + .err() + .unwrap(); + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + assert_eq!(error.status_code(), StatusCode::RateLimited); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + num_rows as u64, + "rejected instant rows are not successful best-effort drops" + ); + assert_eq!( + MAX_MIRROR_PENDING_ROWS as u64, + pending.load(Ordering::Relaxed) + ); + assert_eq!( + gauge_before, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() + ); + } + + /// A source table mapped to several flownodes clones its requests per peer, + /// so the reservation must cover every clone and each spawned task must hand + /// back exactly its own share; otherwise the shared budget drifts. + #[tokio::test] + async fn flow_mirror_reserves_and_releases_per_peer_shares() { + let _guard = lock_mirror_metrics().await; + let per_peer_rows = 10; + let total_rows = (2 * per_peer_rows) as u64; + // Exactly enough room for the whole batch, cloned rows included. + let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows; + let pending = Arc::new(AtomicU64::new(start)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler { + gate: gate.clone(), + affected_rows: per_peer_rows as u64, + })); + + let mut requests = mirror_requests(&flownode_peer(), per_peer_rows); + requests.extend(mirror_requests( + &Peer { + id: 2, + addr: "127.0.0.1:4002".to_string(), + }, + per_peer_rows, + )); + + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + FlowMirrorTask { requests } + .detach(node_manager, pending.clone(), false) + .unwrap(); + + // Both per-peer payloads are counted, so the budget is exactly full + // instead of over-reserved by the duplicated rows. + assert_eq!( + MAX_MIRROR_PENDING_ROWS as u64, + pending.load(Ordering::Relaxed) + ); + + // Let both tasks finish and check that the budget lands back on its + // starting value: over-releasing a shared reservation would saturate it. + // The gauge is checked too because `release_mirror_pending_rows` moves + // the budget before the gauge; exiting on the budget alone could leak a + // pending `gauge.sub` into the next test holding the metrics lock. + gate.add_permits(2); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != start + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "the mirror tasks must release exactly their own shares" + ); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + } + + #[tokio::test] + async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() { + let _guard = lock_mirror_metrics().await; + let peer1 = flownode_peer(); + let peer2 = Peer { + id: 2, + addr: "127.0.0.1:4002".to_string(), + }; + let gate1 = Arc::new(tokio::sync::Semaphore::new(0)); + let gate2 = Arc::new(tokio::sync::Semaphore::new(0)); + let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([ + (peer1.id, gate1.clone()), + (peer2.id, gate2.clone()), + ]))); + let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel(); + let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel(); + let node_manager = Arc::new(RecordingNodeManager { + datanode_dispatch: datanode_tx, + flownode_dispatch: flownode_tx, + gates, + fail_flownode: false, + }); + let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20)); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + let mut requests = mirror_requests(&peer1, 10); + requests.extend(mirror_requests(&peer2, 10)); + let task = FlowMirrorTask { requests }; + task.detach(node_manager.clone(), pending.clone(), false) + .unwrap(); + let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv()) + .await + .unwrap() + .unwrap(); + let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_ne!(first.id, second.id); + let fast_gate = if first.id == peer1.id { + gate1.clone() + } else { + gate2.clone() + }; + let slow_gate = if first.id == peer1.id { + gate2.clone() + } else { + gate1.clone() + }; + fast_gate.add_permits(1); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 { + assert!( + std::time::Instant::now() < deadline, + "fast peer did not release" + ); + std::thread::sleep(Duration::from_millis(1)); + } + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + FlowMirrorTask { + requests: mirror_requests(&peer1, 11), + } + .detach(node_manager, pending.clone(), false) + .unwrap(); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + 11 + ); + assert!(flownode_rx.try_recv().is_err()); + slow_gate.add_permits(1); + while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20 + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "slow peer did not release" + ); + std::thread::sleep(Duration::from_millis(1)); + } + } + + #[tokio::test] + async fn flow_mirror_failure_releases_budget_and_gauge() { + let _guard = lock_mirror_metrics().await; + let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel(); + let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel(); + let node_manager = Arc::new(RecordingNodeManager { + datanode_dispatch: datanode_tx, + flownode_dispatch: flownode_tx, + gates: Arc::new(std::sync::Mutex::new(HashMap::new())), + fail_flownode: true, + }); + let pending = Arc::new(AtomicU64::new(0)); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), 5), + } + .detach(node_manager, pending.clone(), false) + .unwrap(); + assert_eq!( + tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv()) + .await + .unwrap() + .unwrap() + .id, + flownode_peer().id + ); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != 0 + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "failed task leaked reservation" + ); + std::thread::sleep(Duration::from_millis(1)); + } + } + + #[tokio::test] + async fn flow_mirror_pending_counter_consistent() { + let _guard = lock_mirror_metrics().await; + let pending = Arc::new(AtomicU64::new(0)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler { + gate: gate.clone(), + affected_rows: 7, + })); + + // A batch without rows short-circuits and leaves the budget untouched. + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + FlowMirrorTask { + requests: HashMap::new(), + } + .detach(node_manager.clone(), pending.clone(), false) + .unwrap(); + assert_eq!(0, pending.load(Ordering::Relaxed)); + assert_eq!( + gauge_before, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + "an empty batch must not touch the pending gauge" + ); + + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + let num_rows = 7; + + FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), num_rows), + } + .detach(node_manager, pending.clone(), false) + .unwrap(); + + // The spawned task is still waiting on the gated flownode, so the + // reservation made by `detach` is the value observable here. + assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed)); + assert_eq!( + num_rows as i64, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before, + "the pending gauge must track the pending rows" + ); + assert_eq!( + dropped_before, + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + "a batch within the limit must not be counted as dropped" + ); + + // Release the in-flight task and wait for it to hand its reservation + // back, so no task is left holding the shared metrics after this test. + gate.add_permits(1); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != 0 + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "the mirror task did not release its pending reservation" + ); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + } + + #[test] + fn mirror_logs_are_rate_limited() { + let log = Arc::new(MirrorLog::new()); + assert_eq!(Some(1), log.claim_report(1_000)); + assert_eq!(None, log.claim_report(1_500)); + assert_eq!(None, log.claim_report(10_999)); + assert_eq!(Some(3), log.claim_report(11_000)); + assert_eq!(None, log.claim_report(11_001)); + assert_eq!(Some(2), log.claim_report(21_000)); + + let concurrent = Arc::new(MirrorLog::new()); + assert_eq!(Some(1), concurrent.claim_report(50_000)); + let barrier = Arc::new(std::sync::Barrier::new(8)); + let threads = (0..8) + .map(|_| { + let log = concurrent.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + log.claim_report(60_000) + }) + }) + .collect::>(); + let winners = threads + .into_iter() + .filter_map(|thread| thread.join().unwrap()) + .collect::>(); + assert_eq!(winners.len(), 1); + let residual = concurrent.events.swap(0, Ordering::Relaxed); + assert_eq!(winners[0] + residual, 8); + } } diff --git a/src/operator/src/metrics.rs b/src/operator/src/metrics.rs index 3e0b2fed2a6..438904ed36c 100644 --- a/src/operator/src/metrics.rs +++ b/src/operator/src/metrics.rs @@ -47,6 +47,11 @@ lazy_static! { "table operator mirror pending rows" ) .unwrap(); + pub static ref DIST_MIRROR_DROPPED_ROW_COUNT: IntCounter = register_int_counter!( + "greptime_operator_dist_mirror_dropped_rows", + "table operator mirror rows dropped because the pending mirror row limit was exceeded" + ) + .unwrap(); pub static ref DIST_DELETE_ROW_COUNT: IntCounterVec = register_int_counter_vec!( "greptime_table_operator_delete_rows", "table operator delete rows", diff --git a/src/operator/src/request.rs b/src/operator/src/request.rs index b46a6b657a6..ce3b3568651 100644 --- a/src/operator/src/request.rs +++ b/src/operator/src/request.rs @@ -18,7 +18,7 @@ use api::helper::to_pb_time_unit; use api::v1::region::region_request::Body as RegionRequestBody; use api::v1::region::{ BuildIndexRequest, CompactRequest, CompactionTimeRange, FlushRequest, RegionRequestHeader, - TruncateRequest, Unflushed, truncate_request, + TruncateRequest, Unflushed, build_index_request, truncate_request, }; use catalog::CatalogManagerRef; use common_catalog::build_db_string; @@ -107,6 +107,30 @@ impl Requester { request: BuildIndexTableRequest, ctx: QueryContextRef, ) -> Result { + if matches!( + request.options, + Some(build_index_request::Options::SeriesIndex(_)) + ) { + let table = self + .catalog_manager + .table( + &request.catalog_name, + &request.schema_name, + &request.table_name, + None, + ) + .await + .context(CatalogSnafu)?; + let table = table.with_context(|| TableNotFoundSnafu { + table_name: common_catalog::format_full_table_name( + &request.catalog_name, + &request.schema_name, + &request.table_name, + ), + })?; + let info = table.table_info(); + ensure_build_series_index_supported(&info.meta.engine, info.is_physical_table())?; + } let partitions = &self .get_table_partition_info( &request.catalog_name, @@ -121,6 +145,7 @@ impl Requester { .map(|partition| { RegionRequestBody::BuildIndex(BuildIndexRequest { region_id: partition.id.into(), + options: request.options, }) }) .collect(); @@ -404,6 +429,16 @@ impl Requester { } } +fn ensure_build_series_index_supported(engine: &str, is_physical_table: bool) -> Result<()> { + ensure!( + engine == METRIC_ENGINE && is_physical_table, + NotSupportedSnafu { + feat: "building series indexes requires a physical metric table", + } + ); + Ok(()) +} + fn ensure_discard_unflushed_supported(engine: &str, is_physical_table: bool) -> Result<()> { ensure!( engine != METRIC_ENGINE || is_physical_table, @@ -422,6 +457,14 @@ mod tests { use super::*; + #[test] + fn test_build_series_index_requires_physical_metric_table() { + ensure_build_series_index_supported(METRIC_ENGINE, true).unwrap(); + assert!(ensure_build_series_index_supported(METRIC_ENGINE, false).is_err()); + assert!(ensure_build_series_index_supported("mito", false).is_err()); + assert!(ensure_build_series_index_supported("mito", true).is_err()); + } + #[test] fn test_to_pb_compaction_time_range_normalizes_mixed_units() { let range = TimestampRange::new( diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 21b82f51a3b..f9727be6c12 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -15,6 +15,8 @@ use std::collections::HashMap; use std::sync::Arc; +use arrow::array::{Array, ArrayRef, AsArray, UInt64Array, make_array}; +use arrow::datatypes::DataType; use client::OutputData; use common_base::readable_size::ReadableSize; use common_datasource::file_format::Format; @@ -22,6 +24,8 @@ use common_datasource::file_format::csv::stream_to_csv; use common_datasource::file_format::json::stream_to_json; use common_datasource::file_format::parquet::stream_to_parquet; use common_datasource::object_store::build_backend_for_write_with_path; +use common_datasource::packed_writer::PackedTableWriter; +use common_datasource::parquet_writer::ParquetFileWriter; use common_query::Output; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; use common_recordbatch::{ @@ -32,16 +36,22 @@ use common_telemetry::{debug, tracing}; use datafusion::datasource::DefaultTableSource; use datafusion_common::TableReference as DfTableReference; use datafusion_expr::LogicalPlanBuilder; +use futures::StreamExt; use object_store::ObjectStore; use session::context::QueryContextRef; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use table::TableRef; use table::requests::CopyTableRequest; use table::table::adapter::DfTableProviderAdapter; use table::table_reference::TableReference; +use tokio_util::sync::CancellationToken; use crate::error::{self, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, Result}; use crate::statement::StatementExecutor; +use crate::statement::export_logical_tables::writers::ExportWriteBudget; +use crate::statement::export_logical_tables::{ + expand_export_batch, map_writer_error, rows_within_budget, +}; // The buffer size should be greater than 5MB (minimum multipart upload size). /// Buffer size to flush data to object stores. @@ -117,6 +127,18 @@ impl StatementExecutor { table: TableRef, req: CopyTableRequest, query_ctx: QueryContextRef, + ) -> Result { + self.copy_captured_table_to_managed(table, req, query_ctx, None, None) + .await + } + + pub(crate) async fn copy_captured_table_to_managed( + &self, + table: TableRef, + req: CopyTableRequest, + query_ctx: QueryContextRef, + managed: Option<(&ExportWriteBudget, &CancellationToken)>, + packed: Option, ) -> Result { let info = table.table_info(); let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name); @@ -165,7 +187,7 @@ impl StatementExecutor { } = &req; debug!("Copy table: {table_id} to location: {location}"); - self.copy_to_file(&format, output, location, connection) + self.copy_to_file_managed(&format, output, location, connection, managed, packed) .await } @@ -176,9 +198,26 @@ impl StatementExecutor { location: &str, connection: &HashMap, ) -> Result { - let output = output - .map_dictionary_to_values() - .context(error::BuildRecordBatchSnafu)?; + self.copy_to_file_managed(format, output, location, connection, None, None) + .await + } + + async fn copy_to_file_managed( + &self, + format: &Format, + output: Output, + location: &str, + connection: &HashMap, + managed: Option<(&ExportWriteBudget, &CancellationToken)>, + packed: Option, + ) -> Result { + let output = if managed.is_none() { + output + .map_dictionary_to_values() + .context(error::BuildRecordBatchSnafu)? + } else { + output + }; let stream = match output.data { OutputData::Stream(stream) => stream, OutputData::RecordBatches(record_batches) => record_batches.as_stream(), @@ -192,7 +231,441 @@ impl StatementExecutor { let filename = backend.object_path.context(error::UnexpectedSnafu { violated: format!("Expected filename, path: {location}"), })?; - self.stream_to_file(stream, format, backend.object_store, &filename) + if let Some((budget, token)) = managed { + stream_to_managed_parquet_with_packed( + stream, + backend.object_store, + &filename, + budget, + token, + packed, + ) .await + } else { + self.stream_to_file(stream, format, backend.object_store, &filename) + .await + } + } +} + +#[cfg(test)] +pub(crate) async fn stream_to_managed_parquet( + stream: SendableRecordBatchStream, + store: ObjectStore, + path: &str, + budget: &ExportWriteBudget, + token: &CancellationToken, +) -> Result { + stream_to_managed_parquet_with_packed(stream, store, path, budget, token, None).await +} + +async fn stream_to_managed_parquet_with_packed( + mut stream: SendableRecordBatchStream, + store: ObjectStore, + path: &str, + budget: &ExportWriteBudget, + token: &CancellationToken, + packed: Option, +) -> Result { + use common_recordbatch::{RecordBatch, map_dictionary_to_values_schema}; + let original = stream.schema(); + let (expanded_schema, expand) = map_dictionary_to_values_schema(original.clone()); + let json_columns = expanded_schema + .column_schemas() + .iter() + .enumerate() + .filter_map(|(index, column)| { + (column.data_type.is_json() && !column.data_type.is_json2()).then_some(index) + }) + .collect::>(); + let (mapped_schema, json) = map_json_type_to_string_schema(expanded_schema.clone()); + let output_schema = if json { + mapped_schema + } else { + expanded_schema.clone() + }; + let packed_destination = packed.is_some(); + let mut writer = if let Some(packed) = packed { + ParquetFileWriter::open_packed( + output_schema.arrow_schema().clone(), + store, + path, + Some( + crate::statement::export_logical_tables::LogicalTableExportLimits::default().writer, + ), + packed, + ) + } else { + ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None).await + } + .context(error::WriteStreamToFileSnafu { path })?; + let mut started = false; + let result = async { + let mut rows = 0; + loop { + let batch = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + batch = stream.next() => batch, + }; + let Some(batch) = batch else { break }; + let batch = batch + .context(error::BuildRecordBatchSnafu)? + .into_df_record_batch(); + let mut offset = 0; + while offset < batch.num_rows() { + // The scan batch remains query-owned. Bound the downstream slice and + // detach its buffers below if it would retain more than its reservation. + let (conversion, retained) = + ExportWriteBudget::conversion_budget(0, batch.num_columns(), usize::MAX)?; + let input = batch.clone(); + let json_columns = json_columns.clone(); + let (len, estimated) = common_runtime::spawn_blocking_global(move || { + rows_within_budget(&input, offset, input.num_rows(), conversion, &json_columns) + }) + .await + .context(error::JoinTaskSnafu)??; + let reservation = retained.saturating_add(estimated.saturating_mul(4)); + let permit = budget.reserve(reservation, token).await?; + let batch = batch.clone(); + let (original, expanded_schema, output_schema) = ( + original.clone(), + expanded_schema.clone(), + output_schema.clone(), + ); + let (batch, len) = common_runtime::spawn_blocking_global(move || { + let mut batch = RecordBatch::from_df_record_batch( + original.clone(), + batch.slice(offset, len), + ); + if expand { + batch = RecordBatch::from_df_record_batch( + expanded_schema.clone(), + expand_export_batch( + &batch.into_df_record_batch(), + expanded_schema.arrow_schema().clone(), + )?, + ); + } + if json { + batch = map_json_type_to_string(batch, &expanded_schema, &output_schema) + .context(error::BuildRecordBatchSnafu)?; + } + let mut batch = batch.into_df_record_batch(); + if batch.get_array_memory_size() > reservation { + let indices = UInt64Array::from_iter_values(0..batch.num_rows() as u64); + batch = arrow::compute::take_record_batch(&batch, &indices) + .context(error::ComputeArrowSnafu)?; + let arrays = batch + .columns() + .iter() + .map(compact_view_buffers) + .collect::>>()?; + batch = arrow::record_batch::RecordBatch::try_new(batch.schema(), arrays) + .context(error::ComputeArrowSnafu)?; + } + ensure!( + batch.get_array_memory_size() <= reservation, + error::LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok::<_, error::Error>((batch, len)) + }) + .await + .context(error::JoinTaskSnafu)??; + started = true; + let write = writer.write(batch, Some(token)).await; + drop(permit); + write.map_err(|error| map_writer_error(error, path))?; + rows += len; + offset += len; + } + } + started = true; + writer + .finish(Some(token)) + .await + .map_err(|error| map_writer_error(error, path))?; + Ok(rows) + } + .await; + if result.is_err() { + token.cancel(); + if (started || packed_destination) + && let Err(error) = writer.abort().await + { + common_telemetry::warn!(error; "Failed to abort ordinary export file"); + } + } + result +} + +// Arrow take copies ordinary buffers but shares view data buffers, including +// views nested inside lists or structs. Reclaim those unselected values too. +fn compact_view_buffers(array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Utf8View => Ok(Arc::new(array.as_string_view().gc())), + DataType::BinaryView => Ok(Arc::new(array.as_binary_view().gc())), + _ => { + let data = array.to_data(); + if data.child_data().is_empty() { + return Ok(array.clone()); + } + let children = data + .child_data() + .iter() + .map(|child| { + compact_view_buffers(&make_array(child.clone())).map(|array| array.to_data()) + }) + .collect::>>()?; + Ok(make_array( + data.into_builder() + .child_data(children) + .build() + .context(error::ComputeArrowSnafu)?, + )) + } + } +} + +#[cfg(test)] +mod tests { + use arrow::array::{ + ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, Int32Array, StringArray, + StringViewArray, + }; + use arrow::datatypes::Int32Type; + use common_recordbatch::{RecordBatch, RecordBatches}; + use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + + use super::*; + + #[tokio::test] + async fn managed_copy_preserves_existing_file_before_sink_open() { + let temp_dir = common_test_util::temp_dir::create_temp_dir("managed_copy_existing_file"); + let store = object_store::secure_fs::SecureFsRoot::open(temp_dir.path()) + .unwrap() + .build_operator(); + let path = "existing.parquet"; + store.write(path, "original").await.unwrap(); + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = arrow::record_batch::RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + let batches = async_stream::stream! { + yield Ok(batch); + yield Err(datafusion::error::DataFusionError::Execution("source failed".into())); + }; + let stream = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new( + schema.arrow_schema().clone(), + batches, + ); + let stream = + common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(Box::pin(stream)) + .unwrap(); + + let result = stream_to_managed_parquet( + Box::pin(stream), + store.clone(), + path, + &ExportWriteBudget::new(1), + &CancellationToken::new(), + ) + .await; + + assert!(result.is_err()); + assert_eq!( + store.read(path).await.unwrap().to_bytes().as_ref(), + b"original" + ); + } + + #[tokio::test] + async fn managed_copy_rechunks_large_backing_buffers() { + let value = "x".repeat(9 * 1024); + for view in [false, true] { + let data_type = if view { + ConcreteDataType::utf8_view_datatype() + } else { + ConcreteDataType::string_datatype() + }; + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", data_type, false, + )])); + let values = vec![value.as_str(); 8192]; + let array: ArrayRef = if view { + Arc::new(StringViewArray::from(values)) + } else { + Arc::new(StringArray::from(values)) + }; + let batch = arrow::record_batch::RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![array], + ) + .unwrap(); + assert!(batch.get_array_memory_size() > 64 * 1024 * 1024); + // Also exercise a small slice that still references the large allocation. + for input in [batch.clone(), batch.slice(200, 3)] { + let rows = input.num_rows(); + let batches = RecordBatches::try_new( + schema.clone(), + vec![RecordBatch::from_df_record_batch(schema.clone(), input)], + ) + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(1); + assert_eq!( + stream_to_managed_parquet( + batches.as_stream(), + store.clone(), + "large.parquet", + &budget, + &CancellationToken::new(), + ) + .await + .unwrap(), + rows + ); + let reader = ParquetRecordBatchReaderBuilder::try_new( + store.read("large.parquet").await.unwrap().to_bytes(), + ) + .unwrap() + .build() + .unwrap(); + let mut actual_rows = 0; + for batch in reader { + let batch = batch.unwrap(); + let strings = arrow::compute::cast(batch.column(0), &DataType::Utf8).unwrap(); + assert!( + strings + .as_string::() + .iter() + .all(|item| item == Some(value.as_str())) + ); + actual_rows += batch.num_rows(); + } + assert_eq!(actual_rows, rows); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + } + } + } + + #[tokio::test] + async fn managed_copy_preserves_dictionary_json_views_and_empty_schema() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "host", + ConcreteDataType::dictionary_datatype( + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype(), + ), + true, + ), + ColumnSchema::new("json", ConcreteDataType::json_datatype(), true), + ColumnSchema::new("text_view", ConcreteDataType::utf8_view_datatype(), true), + ColumnSchema::new( + "binary_view", + ConcreteDataType::binary_view_datatype(), + true, + ), + ])); + let dictionary = DictionaryArray::::new( + Int32Array::from(vec![Some(0), None, Some(1), Some(1)]), + Arc::new(StringArray::from(vec!["tag", &"x".repeat(3 * 1024 * 1024)])), + ); + let json = datatypes::types::parse_string_to_jsonb(&format!( + r#"{{"value":"escaped\ntext{}","n":123456789}}"#, + "y".repeat(160_000) + )) + .unwrap(); + let arrays = vec![ + Arc::new(dictionary) as ArrayRef, + Arc::new(BinaryArray::from(vec![ + Some(json.as_slice()), + None, + None, + Some(json.as_slice()), + ])), + Arc::new(StringViewArray::from(vec![ + Some("short"), + None, + Some("long view"), + Some("long view"), + ])), + Arc::new(BinaryViewArray::from(vec![ + Some(&b"short"[..]), + None, + Some(&b"long view"[..]), + Some(&b"long view"[..]), + ])), + ]; + let batch = + arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), arrays) + .unwrap(); + let batch = RecordBatch::from_df_record_batch(schema.clone(), batch); + for empty in [false, true] { + let batches = RecordBatches::try_new( + schema.clone(), + if empty { vec![] } else { vec![batch.clone()] }, + ) + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let managed = stream_to_managed_parquet( + batches.as_stream(), + store.clone(), + "managed.parquet", + &budget, + &token, + ) + .await + .unwrap(); + let output = Output::new_with_record_batches(batches) + .map_dictionary_to_values() + .unwrap(); + let OutputData::RecordBatches(batches) = output.data else { + panic!("expected batches"); + }; + let stream = SendableRecordBatchMapper::new( + batches.as_stream(), + map_json_type_to_string, + map_json_type_to_string_schema, + ); + let ordinary = stream_to_parquet( + Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))), + store.clone(), + "ordinary.parquet", + WRITE_CONCURRENCY, + ) + .await + .unwrap(); + assert_eq!(managed, ordinary); + let mut outputs = Vec::new(); + for name in ["managed.parquet", "ordinary.parquet"] { + let reader = ParquetRecordBatchReaderBuilder::try_new( + store.read(name).await.unwrap().to_bytes(), + ) + .unwrap(); + let schema = reader.schema().clone(); + let values = reader + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + outputs.push((schema, values)); + } + assert_eq!(outputs[0], outputs[1]); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + } } } diff --git a/src/operator/src/statement/database_copy.rs b/src/operator/src/statement/database_copy.rs index 117104b9d8f..08162c5d5e6 100644 --- a/src/operator/src/statement/database_copy.rs +++ b/src/operator/src/statement/database_copy.rs @@ -27,6 +27,7 @@ use table::TableRef; use table::metadata::TableType; use table::requests::CopyDatabaseRequest; use table::table_reference::TableReference; +use tokio::sync::Semaphore; use url::Url; use crate::error::{self, Result}; @@ -56,10 +57,10 @@ pub(crate) fn parse_parallelism_from_option_map(options: &HashMap().ok()) .unwrap_or_else(get_total_cpu_cores) - .max(1) + .clamp(1, Semaphore::MAX_PERMITS) } -/// Rejects import-only layouts before either database export path creates output. +/// Rejects layouts unsupported by the per-table export path. pub(crate) fn validate_database_export_layout(options: &HashMap) -> Result<()> { if let Some(layout) = options.get("metric_data_layout") { return error::InvalidCopyParameterSnafu { @@ -321,5 +322,11 @@ mod tests { let options = HashMap::from([("parallelism".to_string(), "0".to_string())]); assert_eq!(parse_parallelism_from_option_map(&options), 1); + + let options = HashMap::from([("parallelism".to_string(), usize::MAX.to_string())]); + assert_eq!( + parse_parallelism_from_option_map(&options), + Semaphore::MAX_PERMITS + ); } } diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 772bff4aad9..fa9da9f322e 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -93,7 +93,9 @@ use sql::statements::create::{ }; use sql::statements::statement::Statement; use sqlparser::ast::{Expr, Ident, UnaryOperator, Value as ParserValue}; -use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; +use store_api::metric_engine_consts::{ + LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, +}; use store_api::mito_engine_options::APPEND_MODE_KEY; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::TableRef; @@ -341,7 +343,22 @@ impl StatementExecutor { #[tracing::instrument(skip_all)] pub async fn create_table(&self, stmt: CreateTable, ctx: QueryContextRef) -> Result { - let (catalog, schema, _table) = table_idents_to_full_name(&stmt.name, &ctx) + let mut create_expr = self.prepare_create_table(&stmt, &ctx).await?; + self.create_table_inner( + &mut create_expr, + stmt.partitions, + ctx, + TriggerReason::Manual, + ) + .await + } + + async fn prepare_create_table( + &self, + stmt: &CreateTable, + ctx: &QueryContextRef, + ) -> Result { + let (catalog, schema, _table) = table_idents_to_full_name(&stmt.name, ctx) .map_err(BoxedError::new) .context(error::ExternalSnafu)?; @@ -356,7 +373,7 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)? .map(|v| v.into_inner()); - let create_expr = &mut expr_helper::create_to_expr(&stmt, &ctx)?; + let mut create_expr = expr_helper::create_to_expr(stmt, ctx)?; // TTL and compaction options are handled separately; ingestion quota is database-only. if let Some(schema_options) = schema_options { for (key, value) in schema_options.extra_options.iter() { @@ -370,8 +387,72 @@ impl StatementExecutor { } } - self.create_table_inner(create_expr, stmt.partitions, ctx, TriggerReason::Manual) - .await + Ok(create_expr) + } + + /// Prepares every explicit logical CREATE before submitting one batch procedure. + pub async fn batch_create_logical_tables( + &self, + statements: Vec, + ctx: QueryContextRef, + ) -> Result<()> { + let mut expressions = Vec::with_capacity(statements.len()); + let mut targets = std::collections::HashSet::new(); + let mut group = None; + for statement in statements { + ensure!( + statement.engine == METRIC_ENGINE_NAME + && statement.options.get(LOGICAL_TABLE_METADATA_KEY).is_some() + && statement + .options + .value(PHYSICAL_TABLE_METADATA_KEY) + .is_none(), + InvalidSqlSnafu { + err_msg: "batch requires explicit Metric logical CREATE TABLE statements" + } + ); + let expr = self.prepare_create_table(&statement, &ctx).await?; + let key = ( + expr.catalog_name.clone(), + expr.schema_name.clone(), + expr.table_options.get(LOGICAL_TABLE_METADATA_KEY).cloned(), + ); + ensure!( + group.as_ref().is_none_or(|group| group == &key), + InvalidSqlSnafu { + err_msg: "batch must share catalog, schema and literal on_physical_table" + } + ); + group = Some(key); + ensure!( + targets.insert(expr.table_name.clone()), + InvalidSqlSnafu { + err_msg: "duplicate target in logical-table batch" + } + ); + self.validate_logical_create(&expr, statement.partitions.as_ref(), &ctx) + .await?; + expressions.push(expr); + } + self.create_logical_tables(&expressions, ctx, TriggerReason::Manual) + .await?; + Ok(()) + } + + async fn validate_logical_create( + &self, + expr: &CreateTableExpr, + partitions: Option<&Partitions>, + ctx: &QueryContextRef, + ) -> Result<()> { + ensure_table_definition_writable(&expr.schema_name, &expr.table_name)?; + if let Some(partitions) = partitions + && !partitions.exprs.is_empty() + { + self.validate_logical_table_partition_rule(expr, partitions, ctx) + .await?; + } + Ok(()) } #[tracing::instrument(skip_all)] @@ -472,12 +553,8 @@ impl StatementExecutor { .table_options .contains_key(LOGICAL_TABLE_METADATA_KEY) { - if let Some(partitions) = partitions.as_ref() - && !partitions.exprs.is_empty() - { - self.validate_logical_table_partition_rule(create_table, partitions, &query_ctx) - .await?; - } + self.validate_logical_create(create_table, partitions.as_ref(), &query_ctx) + .await?; // Create logical tables self.create_logical_tables( std::slice::from_ref(create_table), @@ -2970,7 +3047,7 @@ mod test { use datafusion::functions_aggregate::expr_fn::count; use datafusion::logical_expr::builder::LogicalTableSource; use datafusion::logical_expr::{LogicalPlanBuilder, col}; - use session::context::{QueryContext, QueryContextBuilder}; + use session::context::QueryContext; use sql::dialect::GreptimeDbDialect; use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; @@ -3967,54 +4044,4 @@ WITH ('repartition.column.hint' = 'host')", .contains("cannot set repartition.column.hint on a table with partition metadata") ); } - - #[tokio::test] - #[ignore = "TODO(ruihang): WIP new partition rule"] - async fn test_parse_partitions() { - common_telemetry::init_default_ut_logging(); - let cases = [ - ( - r" -CREATE TABLE rcx ( a INT, b STRING, c TIMESTAMP, TIME INDEX (c) ) -PARTITION ON COLUMNS (b) ( - b < 'hz', - b >= 'hz' AND b < 'sh', - b >= 'sh' -) -ENGINE=mito", - r#"[{"column_list":["b"],"value_list":["{\"Value\":{\"String\":\"hz\"}}"]},{"column_list":["b"],"value_list":["{\"Value\":{\"String\":\"sh\"}}"]},{"column_list":["b"],"value_list":["\"MaxValue\""]}]"#, - ), - ( - r" -CREATE TABLE rcx ( a INT, b STRING, c TIMESTAMP, TIME INDEX (c) ) -PARTITION BY RANGE COLUMNS (b, a) ( - PARTITION r0 VALUES LESS THAN ('hz', 10), - b < 'hz' AND a < 10, - b >= 'hz' AND b < 'sh' AND a >= 10 AND a < 20, - b >= 'sh' AND a >= 20 -) -ENGINE=mito", - r#"[{"column_list":["b","a"],"value_list":["{\"Value\":{\"String\":\"hz\"}}","{\"Value\":{\"Int32\":10}}"]},{"column_list":["b","a"],"value_list":["{\"Value\":{\"String\":\"sh\"}}","{\"Value\":{\"Int32\":20}}"]},{"column_list":["b","a"],"value_list":["\"MaxValue\"","\"MaxValue\""]}]"#, - ), - ]; - let ctx = QueryContextBuilder::default().build().into(); - for (sql, expected) in cases { - let result = ParserContext::create_with_dialect( - sql, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .unwrap(); - match &result[0] { - Statement::CreateTable(c) => { - let expr = expr_helper::create_to_expr(c, &QueryContext::arc()).unwrap(); - let (partitions, _) = - parse_partitions(&expr, c.partitions.clone(), &ctx).unwrap(); - let json = serde_json::to_string(&partitions).unwrap(); - assert_eq!(json, expected); - } - _ => unreachable!(), - } - } - } } diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index 7c5de554487..f363c73d7a2 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -20,6 +20,7 @@ use std::future::Future; use common_datasource::file_format::Format; use common_datasource::object_store::{FILE_SCHEMA, FS_SCHEMA, build_backend_for_write, parse_url}; +use common_datasource::packed_writer::{PackedTableWriter, PackedWriter}; use common_meta::key::table_route::TableRouteValue; use futures::StreamExt; use futures::stream::FuturesUnordered; @@ -37,6 +38,7 @@ use crate::statement::database_copy::{ DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory, validate_database_export_layout, }; +use crate::statement::export_logical_tables::writers::{ExportWriteBudget, retain_error}; use crate::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits}; /// A validated request-scoped selection, not a metadata snapshot or an ACL token. @@ -69,7 +71,13 @@ impl StatementExecutor { req: CopyDatabaseRequest, tables: Vec, ) -> Result { - validate_database_export_layout(&req.with)?; + if req + .with + .get("metric_data_layout") + .is_none_or(|v| v != "packed") + { + validate_database_export_layout(&req.with)?; + } validate_database_directory(&req.location)?; let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?; ensure!( @@ -202,48 +210,111 @@ impl StatementExecutor { } output_files.sort(); let req = &plan.request; - let rows = run_database_export_jobs( - plan.jobs, - parse_parallelism_from_option_map(&req.with), - cancellation, - |job, token| { - let ctx = ctx.clone(); - async move { - match job { - DatabaseExportJob::Metric(unit) => self - .export_logical_tables( - &unit, - &req.location, - &req.connection, - req.time_range.as_ref(), - LogicalTableExportLimits::default(), - &token, - ctx, - ) - .await - .map(|summary| summary.rows), - DatabaseExportJob::Ordinary { table, output } => { - let info = table.table_info(); - let copy = CopyTableRequest { - catalog_name: info.catalog_name.clone(), - schema_name: info.schema_name.clone(), - table_name: info.name.clone(), - location: output.location, - with: req.with.clone(), - connection: req.connection.clone(), - pattern: None, - direction: CopyDirection::Export, - timestamp_range: req.time_range, - limit: None, - }; - self.copy_captured_table_to(table, copy, ctx).await - } + let parallelism = parse_parallelism_from_option_map(&req.with); + let budget = ExportWriteBudget::new(parallelism); + let packed = if req + .with + .get("metric_data_layout") + .is_some_and(|v| v == "packed") + { + let store = + build_backend_for_write(&req.location, &req.connection, &self.local_file_access) + .await + .context(error::BuildBackendSnafu)?; + Some( + PackedWriter::new(store).context(error::WriteStreamToFileSnafu { + path: &req.location, + })?, + ) + } else { + None + }; + let rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| { + let ctx = ctx.clone(); + let budget = budget.clone(); + let packed = packed.clone(); + async move { + match job { + DatabaseExportJob::Metric(unit) => self + .export_logical_tables_managed( + &unit, + &req.location, + &req.connection, + req.time_range.as_ref(), + LogicalTableExportLimits::default(), + &token, + ctx, + budget, + packed, + ) + .await + .map(|summary| summary.rows), + DatabaseExportJob::Ordinary { table, output } => { + let info = table.table_info(); + let copy = CopyTableRequest { + catalog_name: info.catalog_name.clone(), + schema_name: info.schema_name.clone(), + table_name: info.name.clone(), + location: output.location, + with: req.with.clone(), + connection: req.connection.clone(), + pattern: None, + direction: CopyDirection::Export, + timestamp_range: req.time_range, + limit: None, + }; + let _permit = budget.writer(&token).await?; + let destination = packed.map(|shared| { + PackedTableWriter::new(shared, info.name.clone(), info.table_id(), true) + }); + self.copy_captured_table_to_managed( + table, + copy, + ctx, + Some((&budget, &token)), + destination, + ) + .await } } - }, - ) - .await?; - Ok(DatabaseExportSummary { rows, output_files }) + } + }) + .await; + if let Some(packed) = packed { + let mut packed = packed.lock().await; + let result = match rows { + Ok(rows) => packed + .finish(cancellation) + .await + .context(error::WriteStreamToFileSnafu { + path: &req.location, + }) + .and_then(|files| { + Ok(DatabaseExportSummary { + rows, + output_files: files + .into_iter() + .map(|file| { + DatabaseExportFile::new(&req.location, &file, "") + .map(|f| f.location) + }) + .collect::>()?, + }) + }), + Err(error) => Err(error), + }; + if result.is_err() + && let Err(error) = packed.abort().await + { + common_telemetry::warn!(error; "Failed to abort Metric pack"); + } + result + } else { + Ok(DatabaseExportSummary { + rows: rows?, + output_files, + }) + } } } @@ -261,6 +332,7 @@ async fn run_database_export_jobs>>( loop { while first_error.is_none() && !cancellation.is_cancelled() + && !token.is_cancelled() && active.len() < parallelism.max(1) { let Some(job) = jobs.next() else { break }; @@ -284,11 +356,12 @@ async fn run_database_export_jobs>>( }; match result { Some(Ok(count)) => rows += count, - Some(Err(err)) if first_error.is_none() => { - first_error = Some(err); + Some(Err(err)) => { + if !cancellation.is_cancelled() || first_error.is_none() { + retain_error(&mut first_error, err); + } token.cancel(); } - Some(Err(err)) => common_telemetry::warn!(err; "Failed to drain database export job"), None => break, } } @@ -339,7 +412,7 @@ mod tests { } .fail(); } - // Ordinary COPY continues its I/O even when Metric jobs cancel. + // Already-started ordinary I/O drains after cancellation. token.cancelled().await; started.send(10).unwrap(); finish_io.acquire().await.unwrap().forget(); diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index 7547a3f54de..368335f8450 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -19,16 +19,22 @@ //! exclusively by the export attempt. Completion and cleanup of closed files //! belong to the caller's chunk protocol, not individual file completion. +pub(crate) mod writers; + use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; -use arrow::array::{Array, AsArray, UInt32Array}; -use arrow::compute::cast; +use arrow::array::{Array, ArrayRef, AsArray, ListArray, StructArray, UInt32Array}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::{cast, take}; use arrow::datatypes::{DataType, SchemaRef}; use arrow::downcast_dictionary_array; use arrow::record_batch::RecordBatch; use common_datasource::object_store::build_backend_for_write; -use common_datasource::parquet_writer::{ParquetFileWriter, ParquetWriterLimits}; +use common_datasource::packed_writer::{PackedTableWriter, PackedWriterRef}; +use common_datasource::parquet_writer::{ + ParquetCreationPolicy, ParquetFileWriter, ParquetWriterLimits, +}; use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; use common_query::OutputData; use common_recordbatch::SendableRecordBatchStream; @@ -52,6 +58,7 @@ use tokio_util::sync::CancellationToken; use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result}; use crate::statement::StatementExecutor; use crate::statement::database_copy::DatabaseExportFile; +use crate::statement::export_logical_tables::writers::{ExportWriteBudget, Payload, TableWriters}; /// Export preprocessing and per-file limits. Query memory and spill remain /// governed by the query engine. @@ -101,7 +108,8 @@ pub struct LogicalTableExport { logical_tables: BTreeMap, } -struct LogicalTableProjection { +pub(crate) struct LogicalTableProjection { + name: String, output: DatabaseExportFile, schema: SchemaRef, projection: Vec, @@ -190,6 +198,7 @@ impl LogicalTableExport { .insert( info.table_id(), LogicalTableProjection { + name: name.clone(), output: DatabaseExportFile::new(directory, name, ".parquet")?, schema, projection: indices, @@ -306,6 +315,33 @@ impl StatementExecutor { limits: LogicalTableExportLimits, cancellation: &CancellationToken, query_ctx: QueryContextRef, + ) -> Result { + self.export_logical_tables_managed( + unit, + directory, + connection, + time_range, + limits, + &cancellation.child_token(), + query_ctx, + ExportWriteBudget::new(1), + None, + ) + .await + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn export_logical_tables_managed( + &self, + unit: &LogicalTableExport, + directory: &str, + connection: &HashMap, + time_range: Option<&TimestampRange>, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + query_ctx: QueryContextRef, + budget: Arc, + packed: Option, ) -> Result { limits.validate()?; let (store, stream) = tokio::select! { @@ -325,10 +361,11 @@ impl StatementExecutor { Ok((store, stream)) } => result?, }; - export_stream(unit, stream, &store, limits, cancellation).await + export_stream_managed(unit, stream, &store, limits, cancellation, budget, packed).await } } +#[cfg(test)] async fn export_stream( unit: &LogicalTableExport, stream: SendableRecordBatchStream, @@ -336,19 +373,45 @@ async fn export_stream( limits: LogicalTableExportLimits, cancellation: &CancellationToken, ) -> Result { - let mut active = None; - let result = write_tables(unit, stream, store, limits, cancellation, &mut active).await; - if result.is_err() - && let Some(writer) = active - && let Err(cleanup_error) = writer - .writer - .abort() - .await - .map_err(|error| map_writer_error(error, &writer.path)) - { - common_telemetry::warn!(cleanup_error; "Failed to clean up incomplete Metric export file"); - } - result + export_stream_managed( + unit, + stream, + store, + limits, + &cancellation.child_token(), + ExportWriteBudget::new(1), + None, + ) + .await +} + +async fn export_stream_managed( + unit: &LogicalTableExport, + stream: SendableRecordBatchStream, + store: &ObjectStore, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + budget: Arc, + packed: Option, +) -> Result { + let mut writers = TableWriters::new(budget.clone()); + writers.packed = packed; + let result = write_tables( + unit, + stream, + store, + limits, + cancellation, + &mut writers, + &budget, + ) + .await; + let (summary, result) = match result { + Ok(summary) => (summary, Ok(())), + Err(error) => (LogicalTableExportSummary::default(), Err(error)), + }; + writers.drain(result, cancellation).await?; + Ok(summary) } fn check_cancelled(cancellation: &CancellationToken) -> Result<()> { @@ -365,7 +428,8 @@ async fn write_tables( store: &ObjectStore, limits: LogicalTableExportLimits, cancellation: &CancellationToken, - active: &mut Option, + writers: &mut TableWriters, + budget: &ExportWriteBudget, ) -> Result { let id_index = stream @@ -423,22 +487,19 @@ async fn write_tables( while end < batch.num_rows() && ids.value(end) == id { end += 1; } - if active.as_ref().is_some_and(|writer| writer.table_id != id) { - finish_active(active, cancellation).await?; + if writers.table_id().is_some_and(|last| last != id) { + writers.close_input(); } check_cancelled(cancellation)?; if let Some(file) = unit.logical_tables.get(&id) { - if active.is_none() { - *active = Some(ActiveWriter::open(id, file, store, limits).await?); + if writers.table_id().is_none() { + writers.open(id, file, store, limits, cancellation).await?; written.insert(id); summary.files += 1; } let projected = batch .project(&file.projection) .context(error::ProjectSchemaSnafu)?; - let writer = active.as_mut().context(error::UnexpectedSnafu { - violated: "missing logical writer", - })?; let mut offset = start; while offset < end { let (expanded, consumed) = expand_bounded_slice( @@ -447,15 +508,12 @@ async fn write_tables( offset, end, limits.conversion_bytes, + budget, + cancellation, ) .await?; check_cancelled(cancellation)?; - // Do not drop an in-flight file operation before cleanup. - writer - .writer - .write(expanded, Some(cancellation)) - .await - .map_err(|error| map_writer_error(error, &writer.path))?; + writers.send(expanded, cancellation).await?; check_cancelled(cancellation)?; offset += consumed; summary.rows += consumed; @@ -466,12 +524,12 @@ async fn write_tables( start = end; } } - finish_active(active, cancellation).await?; + writers.close_input(); for (&id, file) in &unit.logical_tables { if !written.contains(&id) { check_cancelled(cancellation)?; - *active = Some(ActiveWriter::open(id, file, store, limits).await?); - finish_active(active, cancellation).await?; + writers.open(id, file, store, limits, cancellation).await?; + writers.close_input(); summary.files += 1; } } @@ -480,85 +538,178 @@ async fn write_tables( } struct ActiveWriter { - table_id: u32, path: String, writer: ParquetFileWriter, } impl ActiveWriter { - async fn open( + fn open_packed( + table: &LogicalTableProjection, id: u32, + store: &ObjectStore, + limits: LogicalTableExportLimits, + shared: PackedWriterRef, + ) -> Result { + let path = format!("table-{id}.parquet"); + let packed = PackedTableWriter::new(shared, table.name.clone(), id, false); + let writer = ParquetFileWriter::open_packed( + table.schema.clone(), + store.clone(), + &path, + Some(limits.writer), + packed, + ) + .map_err(|e| map_writer_error(e, &path))?; + Ok(Self { path, writer }) + } + + async fn open( table: &LogicalTableProjection, store: &ObjectStore, limits: LogicalTableExportLimits, ) -> Result { let path = table.output.path.clone(); - ensure!( - !store - .exists(&path) - .await - .context(error::ReadObjectSnafu { path: &path })?, - InvalidLogicalTableExportSnafu { - reason: format!("output already exists: {path}") - } - ); - let writer = ParquetFileWriter::open( + let conditional = store.info().capability().write_with_if_not_exists; + if !conditional { + ensure!( + !store + .exists(&path) + .await + .context(error::ReadObjectSnafu { path: &path })?, + InvalidLogicalTableExportSnafu { + reason: format!("output already exists: {path}") + } + ); + } + let writer = ParquetFileWriter::open_with_creation( table.schema.clone(), store.clone(), &path, 1, Some(limits.writer), + if conditional { + ParquetCreationPolicy::IfNotExists + } else { + ParquetCreationPolicy::Overwrite + }, ) .await .map_err(|error| map_writer_error(error, &path))?; - Ok(Self { - table_id: id, - path, - writer, - }) + Ok(Self { path, writer }) } } -async fn finish_active( - active: &mut Option, - cancellation: &CancellationToken, -) -> Result<()> { - if let Some(writer) = active.as_mut() { - writer - .writer - .finish(Some(cancellation)) - .await - .map_err(|error| map_writer_error(error, &writer.path))?; - check_cancelled(cancellation)?; - *active = None; - } - Ok(()) -} - async fn expand_bounded_slice( batch: RecordBatch, schema: SchemaRef, start: usize, end: usize, - budget: usize, -) -> Result<(RecordBatch, usize)> { + requested: usize, + budget: &ExportWriteBudget, + cancellation: &CancellationToken, +) -> Result<(Payload, usize)> { + let (conversion, retained) = ExportWriteBudget::conversion_budget( + batch.get_array_memory_size(), + batch.num_columns(), + requested, + )?; + let input = batch.clone(); + let (len, estimated) = common_runtime::spawn_blocking_global(move || { + rows_within_budget(&input, start, end, conversion, &[]) + }) + .await + .context(error::JoinTaskSnafu)??; + let reservation = retained.saturating_add(estimated.saturating_mul(4)); + let permit = budget.reserve(reservation, cancellation).await?; common_runtime::spawn_blocking_global(move || { - let len = rows_within_budget(&batch, start, end, budget)?; let slice = batch.slice(start, len); - let arrays = slice - .columns() - .iter() - .zip(schema.fields()) - .map(|(array, field)| cast(array, field.data_type()).context(error::ComputeArrowSnafu)) - .collect::>>()?; - let expanded = RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)?; - Ok((expanded, len)) + let expanded = expand_export_batch(&slice, schema)?; + ensure!( + expanded.get_array_memory_size() <= reservation, + LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok(( + Payload { + batch: expanded, + permit, + }, + len, + )) }) .await .context(error::JoinTaskSnafu)? } -fn map_writer_error(source: common_datasource::error::Error, path: &str) -> error::Error { +/// Limits nested casts to selected children: Arrow's List cast otherwise expands +/// the entire values array, even when the parent has been sliced. +pub(crate) fn expand_export_batch(batch: &RecordBatch, schema: SchemaRef) -> Result { + let arrays = batch + .columns() + .iter() + .zip(schema.fields()) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu) +} + +fn expand_export_array(array: &ArrayRef, target: &DataType) -> Result { + if array.data_type() == target { + return Ok(array.clone()); + } + match (array.data_type(), target) { + (DataType::Dictionary(_, _), _) => { + downcast_dictionary_array! { + array => { + // Select dictionary entries before recursively converting nested values. + let selected = take(array.values().as_ref(), array.keys(), None) + .context(error::ComputeArrowSnafu)?; + expand_export_array(&selected, target) + }, + _ => error::UnexpectedSnafu { violated: "invalid dictionary array" }.fail(), + } + } + (DataType::List(_), DataType::List(field)) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let start = offsets[0]; + let end = offsets[list.len()]; + let values = list.values().slice(start as usize, (end - start) as usize); + let values = expand_export_array(&values, field.data_type())?; + let offsets = OffsetBuffer::new( + offsets + .iter() + .map(|offset| offset - start) + .collect::>() + .into(), + ); + Ok(Arc::new( + ListArray::try_new(field.clone(), offsets, values, list.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + (DataType::Struct(_), DataType::Struct(fields)) => { + let array = array.as_struct(); + let columns = array + .columns() + .iter() + .zip(fields) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + Ok(Arc::new( + StructArray::try_new(fields.clone(), columns, array.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + _ => cast(array, target).context(error::ComputeArrowSnafu), + } +} + +pub(crate) fn map_writer_error( + source: common_datasource::error::Error, + path: &str, +) -> error::Error { match source { common_datasource::error::Error::ParquetWriteCancelled {} => { error::LogicalTableExportCancelledSnafu.build() @@ -572,6 +723,12 @@ fn map_writer_error(source: common_datasource::error::Error, path: &str) -> erro common_datasource::error::Error::ParquetWriterResource { reason } => { LogicalTableExportResourceSnafu { reason }.build() } + common_datasource::error::Error::WriteObject { error, .. } + if error.kind() == object_store::ErrorKind::ConditionNotMatch => + { + let reason = format!("output already exists: {path}"); + InvalidLogicalTableExportSnafu { reason }.build() + } source => error::WriteStreamToFileSnafu { path }.into_error(source), } } @@ -579,7 +736,7 @@ fn map_writer_error(source: common_datasource::error::Error, path: &str) -> erro // Count only selected logical values before dictionary expansion, including nested // histogram lists/structs. Offset and validity overhead is charged per value. fn estimate_value_size(array: &dyn Array, row: usize) -> Result { - if array.is_null(row) { + if array.is_null(row) && !matches!(array.data_type(), DataType::Struct(_) | DataType::List(_)) { return Ok(32); } let bytes = match array.data_type() { @@ -587,8 +744,10 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result { DataType::Null => 0, DataType::Utf8 => array.as_string::().value(row).len(), DataType::LargeUtf8 => array.as_string::().value(row).len(), + DataType::Utf8View => array.as_string_view().value(row).len(), DataType::Binary => array.as_binary::().value(row).len(), DataType::LargeBinary => array.as_binary::().value(row).len(), + DataType::BinaryView => array.as_binary_view().value(row).len(), DataType::Struct(_) => { array .as_struct() @@ -629,18 +788,31 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result { Ok(bytes.saturating_add(16)) } -fn rows_within_budget( +pub(crate) fn rows_within_budget( batch: &RecordBatch, start: usize, end: usize, budget: usize, -) -> Result { + json_columns: &[usize], +) -> Result<(usize, usize)> { let mut bytes = 0usize; let mut row = start; while row < end { - let row_bytes = batch.columns().iter().try_fold(0usize, |sum, array| { - Ok::<_, error::Error>(sum.saturating_add(estimate_value_size(array.as_ref(), row)?)) - })?; + let row_bytes = + batch + .columns() + .iter() + .enumerate() + .try_fold(0usize, |sum, (index, array)| { + // JSON escaping and number formatting only expand JSON values. + let expansion = if json_columns.contains(&index) && !array.is_null(row) { + 8 + } else { + 1 + }; + let size = estimate_value_size(array.as_ref(), row)?; + Ok::<_, error::Error>(sum.saturating_add(size.saturating_mul(expansion))) + })?; if row_bytes > budget.saturating_sub(bytes) { break; } @@ -653,7 +825,7 @@ fn rows_within_budget( reason: "one expanded logical row exceeds conversion byte budget" } ); - Ok(row - start) + Ok((row - start, bytes)) } #[cfg(test)] diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 4576c1e395d..7aa6daab127 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -18,6 +18,7 @@ use arrow::array::{ use arrow::buffer::OffsetBuffer; use arrow::datatypes::{Field, Schema, UInt32Type}; use bytes::Bytes; +use common_error::status_code::StatusCode; use common_recordbatch::{RecordBatch as GreptimeRecordBatch, RecordBatches}; use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use table::test_util::EmptyTable; @@ -137,13 +138,15 @@ async fn routes_across_batches_and_writes_empty_files() { vec![Some("a"), Some("b"), Some("unselected")], ), ]; - let result = write_tables( + let budget = ExportWriteBudget::new(4); + let result = export_stream_managed( &unit, stream(batches), &store, export_limits(), &CancellationToken::new(), - &mut None, + budget.clone(), + None, ) .await .unwrap(); @@ -174,6 +177,7 @@ async fn routes_across_batches_and_writes_empty_files() { let (schema, requests) = read(&store, "requests.parquet").await; assert_eq!(schema.fields(), unit.logical_tables[&1027].schema.fields()); assert_eq!(requests[0].column(1).null_count(), 1); + assert_eq!(budget.available(), (4, 64 * 1024 * 1024)); } #[tokio::test] @@ -211,21 +215,16 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() { ), ] { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); - let mut active = None; - let err = write_tables( + let err = export_stream( &unit, stream(batches), &store, limits, &CancellationToken::new(), - &mut active, ) .await .unwrap_err(); assert!(err.to_string().contains(message), "{err}"); - if let Some(writer) = active { - writer.writer.abort().await.unwrap(); - } } } @@ -233,17 +232,17 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() { async fn existing_outputs_are_not_overwritten() { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); store.write("cpu.v1.parquet", "keep").await.unwrap(); - let err = write_tables( + let err = export_stream( &unit(), stream(vec![batch(vec![Some(1025)], vec![None])]), &store, LogicalTableExportLimits::default(), &CancellationToken::new(), - &mut None, ) .await .unwrap_err(); - assert!(err.to_string().contains("already exists")); + let status = common_error::ext::ErrorExt::status_code(&err); + assert_eq!(status, StatusCode::InvalidArguments); assert_eq!( store.read("cpu.v1.parquet").await.unwrap().to_bytes(), Bytes::from_static(b"keep") @@ -272,10 +271,10 @@ fn dictionary_and_nested_histogram_values_are_bounded_before_expansion() { ("histogram", Arc::new(histogram)), ]) .unwrap(); - assert_eq!(rows_within_budget(&batch, 0, 3, 4300).unwrap(), 1); + assert_eq!(rows_within_budget(&batch, 0, 3, 4300, &[]).unwrap().0, 1); // The dictionary and container overhead fit; the nested list elements do not. - assert!(rows_within_budget(&batch, 0, 3, 4180).is_err()); - assert_eq!(rows_within_budget(&batch, 0, 3, 15000).unwrap(), 3); + assert!(rows_within_budget(&batch, 0, 3, 4180, &[]).is_err()); + assert_eq!(rows_within_budget(&batch, 0, 3, 15000, &[]).unwrap().0, 3); } #[test] @@ -402,12 +401,17 @@ async fn native_histogram_parquet_roundtrip() { struct PausedFileWriter { inner: Option, + paused: bool, started: Arc, release: Arc, } impl object_store::layers::mock::oio::Write for PausedFileWriter { async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + if self.paused { + return self.inner.as_mut().unwrap().write(bytes).await; + } + self.paused = true; let mut inner = self.inner.take().unwrap(); let started = self.started.clone(); let release = self.release.clone(); @@ -440,64 +444,216 @@ impl object_store::layers::mock::oio::Write for PausedFileWriter { } #[tokio::test] -async fn cancellation_waits_for_file_creation_before_cleanup() { - use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; - let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); - let access = - common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); - let store = build_backend_for_write( - &format!("{}/", directory.path().display()), - &HashMap::new(), - &access, - ) - .await - .unwrap(); - let started = Arc::new(tokio::sync::Notify::new()); - let release = Arc::new(tokio::sync::Notify::new()); - let factory: MockWriterFactory = Arc::new({ - let started = started.clone(); - let release = release.clone(); - move |_, _, inner| { - Box::new(PausedFileWriter { - inner: Some(inner), - started: started.clone(), - release: release.clone(), - }) +async fn cancellation_drains_storage_and_preserves_committed_files() { + for (large, packed) in [(false, false), (true, false), (false, true)] { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); + let access = + common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); + let store = build_backend_for_write( + &format!("{}/", directory.path().display()), + &HashMap::new(), + &access, + ) + .await + .unwrap(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let factory: MockWriterFactory = Arc::new({ + let started = started.clone(); + let release = release.clone(); + move |_, _, inner| { + Box::new(PausedFileWriter { + inner: Some(inner), + paused: false, + started: started.clone(), + release: release.clone(), + }) + } + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let store = store.layer(object_store::layers::CapabilityOverrideLayer::new( + move |mut capability| { + if packed { + capability.write_multi_max_size = Some(256); + } + capability + }, + )); + let packed = packed + .then(|| common_datasource::packed_writer::PackedWriter::new(store.clone()).unwrap()); + let cancellation = CancellationToken::new(); + let (unit, input) = if large { + let field = Field::new("value", DataType::Utf8, true); + let unit = LogicalTableExport::try_new( + table( + 1024, + "phy", + vec![Field::new(TABLE_ID, DataType::UInt32, false), field.clone()], + true, + ), + &[table(1025, "cpu.v1", vec![field], false)], + ) + .unwrap(); + let mut state = 17u64; + let strings = (0..512) + .map(|_| { + (0..32768) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + char::from(b' ' + (state % 95) as u8) + }) + .collect::() + }) + .collect::>(); + let input = RecordBatch::try_from_iter([ + ( + TABLE_ID, + Arc::new(UInt32Array::from(vec![1025; 512])) as ArrayRef, + ), + ("value", Arc::new(StringArray::from(strings))), + ]) + .unwrap(); + (unit, input) + } else { + (unit(), batch(vec![Some(1025)], vec![Some("a")])) + }; + let budget = ExportWriteBudget::new(1); + let mut limits = LogicalTableExportLimits::default(); + limits.writer.flush_threshold_bytes = 1024 * 1024; + let export = export_stream_managed( + &unit, + stream(vec![input]), + &store, + limits, + &cancellation, + budget.clone(), + packed.clone(), + ); + tokio::pin!(export); + tokio::select! { + result = &mut export => panic!("export completed before the file open: {result:?}"), + _ = started.notified() => {}, } - }); - let store = store.layer( - MockLayerBuilder::default() - .writer_factory(factory) - .build() - .unwrap(), - ); - let cancellation = CancellationToken::new(); - let unit = unit(); - let export = export_stream( - &unit, - stream(vec![batch(vec![Some(1025)], vec![Some("a")])]), - &store, - LogicalTableExportLimits::default(), - &cancellation, - ); - tokio::pin!(export); - tokio::select! { - result = &mut export => panic!("export completed before the file open: {result:?}"), - _ = started.notified() => {}, + if large { + assert!(budget.available().1 < 64 * 1024 * 1024); + } + let held = budget.available(); + assert_eq!(held.0, 0); + cancellation.cancel(); + assert!(futures::poll!(&mut export).is_pending()); + assert_eq!(budget.available().0, held.0); + if large { + // Cancelled admission releases its pending reservation; the worker + // still owns payload capacity while the storage operation is paused. + assert!(budget.available().1 < 64 * 1024 * 1024); + } + release.notify_one(); + let result = export.await; + assert!(matches!( + result, + Err(error::Error::LogicalTableExportCancelled { .. }) + )); + if let Some(packed) = &packed { + packed.lock().await.abort().await.unwrap(); + assert!(!store.exists("pack-000000.bin").await.unwrap()); + assert!(!store.exists("pack-index.json").await.unwrap()); + } + assert_eq!( + store.exists("cpu.v1.parquet").await.unwrap(), + !large && packed.is_none() + ); + if !large && packed.is_none() { + let (_, batches) = read(&store, "cpu.v1.parquet").await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); } - cancellation.cancel(); - assert!(futures::poll!(&mut export).is_pending()); - release.notify_one(); - let result = export.await; - assert!(matches!( - result, - Err(error::Error::LogicalTableExportCancelled { .. }) - )); - assert!(!store.exists("cpu.v1.parquet").await.unwrap()); } struct FailedAbortWriter(object_store::layers::mock::oio::Writer); +#[tokio::test] +async fn metric_fallback_and_ordinary_preserve_ambiguous_commits() { + use object_store::layers::CapabilityOverrideLayer; + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + + use crate::statement::copy_table_to::stream_to_managed_parquet; + + struct AmbiguousCommit(oio::Writer); + impl oio::Write for AmbiguousCommit { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await?; + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "lost close reply", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + Err(object_store::Error::new( + object_store::ErrorKind::Unsupported, + "cannot abort", + )) + } + } + let factory: MockWriterFactory = Arc::new(|_, args, inner| { + assert!(!args.if_not_exists()); + Box::new(AmbiguousCommit(inner)) + }); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer(CapabilityOverrideLayer::new(|mut capability| { + capability.write_with_if_not_exists = false; + capability + })) + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + assert!(!store.info().capability().write_with_if_not_exists); + let rows = || stream(vec![batch(vec![Some(1025)], vec![Some("a")])]); + assert!( + export_stream( + &unit(), + rows(), + &store, + export_limits(), + &CancellationToken::new() + ) + .await + .is_err() + ); + let budget = ExportWriteBudget::new(1); + assert!( + stream_to_managed_parquet( + rows(), + store.clone(), + "ordinary.parquet", + &budget, + &CancellationToken::new() + ) + .await + .is_err() + ); + for path in ["cpu.v1.parquet", "ordinary.parquet"] { + let (_, batches) = read(&store, path).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); +} + impl object_store::layers::mock::oio::Write for FailedAbortWriter { async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { self.0.write(bytes).await @@ -574,3 +730,245 @@ async fn validates_membership_by_table_route() { } } } + +#[tokio::test] +async fn retained_backing_is_reserved_before_conversion_and_until_payload_drop() { + let input = batch(vec![Some(1025); 8192], vec![Some("shared"); 8192]); + let tiny = input.slice(0, 1); + let schema = Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])); + let projected = tiny.project(&[2]).unwrap(); + let full = input.project(&[2]).unwrap().get_array_memory_size(); + assert_eq!(projected.get_array_memory_size(), full); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let blocker = budget.reserve(64 * 1024 * 1024, &token).await.unwrap(); + let convert = expand_bounded_slice(projected, schema, 0, 1, 1024, &budget, &token); + tokio::pin!(convert); + assert!(futures::poll!(&mut convert).is_pending()); + drop(blocker); + let (payload, rows) = convert.await.unwrap(); + assert_eq!(rows, 1); + assert!( + 64 * 1024 * 1024 - budget.available().1 >= full + payload.batch.get_array_memory_size() + ); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + assert!(budget.reserve(64 * 1024 * 1024 + 1, &token).await.is_err()); +} + +#[tokio::test] +async fn groups_and_ordinary_files_share_writer_admission_and_drain() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + use tokio::sync::Semaphore; + + use crate::statement::copy_table_to::stream_to_managed_parquet; + + struct CountedWriter { + inner: oio::Writer, + active: Arc, + started: Arc, + release: Arc, + } + impl Drop for CountedWriter { + fn drop(&mut self) { + self.active.fetch_sub(1, Ordering::SeqCst); + } + } + impl oio::Write for CountedWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.inner.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.started.add_permits(1); + self.release.acquire().await.unwrap().forget(); + self.inner.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.abort().await + } + } + for parallelism in [1, 4] { + for cancel in [false, true] { + let active = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Semaphore::new(0)); + let release = Arc::new(Semaphore::new(0)); + let factory: MockWriterFactory = Arc::new({ + let (active, peak, started, release) = ( + active.clone(), + peak.clone(), + started.clone(), + release.clone(), + ); + move |_, args, inner| { + assert_eq!(args.concurrent(), 1); + let now = active.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(now, Ordering::SeqCst); + Box::new(CountedWriter { + inner, + active: active.clone(), + started: started.clone(), + release: release.clone(), + }) + } + }); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut a = unit(); + let mut b = unit(); + for (group, unit) in [("a", &mut a), ("b", &mut b)] { + for table in unit.logical_tables.values_mut() { + table.output.path = format!("{group}/{}", table.output.path); + } + } + let rows = || { + stream(vec![batch( + vec![Some(1025), Some(1027)], + vec![Some("a"), Some("b")], + )]) + }; + let ordinary = async { + let _permit = budget.writer(&token).await?; + stream_to_managed_parquet( + rows(), + store.clone(), + "ordinary.parquet", + &budget, + &token, + ) + .await + }; + let work = async { + tokio::join!( + export_stream_managed( + &a, + rows(), + &store, + export_limits(), + &token, + budget.clone(), + None, + ), + export_stream_managed( + &b, + rows(), + &store, + export_limits(), + &token, + budget.clone(), + None, + ), + ordinary, + ) + }; + tokio::pin!(work); + tokio::select! { + _ = started.acquire_many(parallelism as u32) => {}, + result = &mut work => panic!("completed while close paused: {result:?}"), + } + assert_eq!(budget.available().0, 0); + assert_eq!(peak.load(Ordering::SeqCst), parallelism); + if cancel { + token.cancel(); + } + assert!(futures::poll!(&mut work).is_pending()); + release.add_permits(10); + let (a, b, ordinary) = work.await; + assert_eq!(a.is_err(), cancel); + assert_eq!(b.is_err(), cancel); + assert_eq!(ordinary.is_err(), cancel); + assert_eq!(active.load(Ordering::SeqCst), 0); + assert!(peak.load(Ordering::SeqCst) <= parallelism); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } + } +} + +#[tokio::test] +async fn nested_dictionary_conversion_limits_child_ranges_and_charges_null_parents() { + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0; 1024]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let list = ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::new(vec![0i32, 1023, 1024].into()), + values, + None, + ); + let input = RecordBatch::try_from_iter([("list", Arc::new(list) as ArrayRef)]).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "list", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + )])); + let budget = ExportWriteBudget::new(1); + let (payload, rows) = expand_bounded_slice( + input, + schema, + 1, + 2, + 8192, + &budget, + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(rows, 1); + assert!(payload.batch.get_array_memory_size() < 16384); + let list = payload.batch.column(0).as_list::(); + assert_eq!(list.value_offsets(), &[0, 1]); + assert_eq!(list.values().as_string::().value(0), "x".repeat(4096)); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let structure = StructArray::new( + vec![Arc::new(Field::new( + "child", + values.data_type().clone(), + true, + ))] + .into(), + vec![values], + Some(arrow::buffer::NullBuffer::from(vec![false])), + ); + let input = RecordBatch::try_from_iter([("struct", Arc::new(structure) as ArrayRef)]).unwrap(); + assert!(rows_within_budget(&input, 0, 1, 128, &[]).is_err()); +} + +#[tokio::test] +async fn completed_table_workers_are_reaped_during_admission() { + for parallelism in [1, 4] { + let mut unit = unit(); + let file = unit.logical_tables.get_mut(&1025).unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut writers = TableWriters::new(budget.clone()); + for id in 0..64 { + file.output.path = format!("empty-{id}.parquet"); + writers + .open(id, file, &store, export_limits(), &token) + .await + .unwrap(); + assert!(writers.pending_tasks() <= parallelism + 1); + } + writers.drain(Ok(()), &token).await.unwrap(); + assert_eq!(writers.pending_tasks(), 0); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } +} diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs new file mode 100644 index 00000000000..70b3bcd2e32 --- /dev/null +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -0,0 +1,317 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use arrow::record_batch::RecordBatch; +use common_datasource::packed_writer::PackedWriterRef; +use futures::stream::FuturesUnordered; +use futures::{FutureExt, StreamExt}; +use object_store::ObjectStore; +use snafu::{OptionExt, ResultExt, ensure}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, Result}; +use crate::statement::export_logical_tables::{ + ActiveWriter, LogicalTableExportLimits, LogicalTableProjection, check_cancelled, + map_writer_error, +}; + +const PAYLOAD_BYTES: usize = 64 * 1024 * 1024; + +/// Request-wide downstream payload and writer admission, separate from query memory. +pub(crate) struct ExportWriteBudget { + writers: Arc, + bytes: Arc, + max_writers: usize, +} + +impl ExportWriteBudget { + pub(crate) fn new(parallelism: usize) -> Arc { + let max_writers = parallelism.max(1); + Arc::new(Self { + writers: Arc::new(Semaphore::new(max_writers)), + bytes: Arc::new(Semaphore::new(PAYLOAD_BYTES)), + max_writers, + }) + } + + #[cfg(test)] + pub(crate) fn available(&self) -> (usize, usize) { + ( + self.writers.available_permits(), + self.bytes.available_permits(), + ) + } + + pub(crate) async fn writer(&self, token: &CancellationToken) -> Result { + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + permit = self.writers.clone().acquire_owned() => permit.map_err(|_| error::UnexpectedSnafu { violated: "writer budget closed" }.build()), + } + } + + pub(crate) async fn reserve( + &self, + size: usize, + token: &CancellationToken, + ) -> Result { + ensure!( + size <= PAYLOAD_BYTES, + error::LogicalTableExportResourceSnafu { + reason: "batch backing buffers exceed request payload budget" + } + ); + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + permit = self.bytes.clone().acquire_many_owned(size as u32) => permit.map_err(|_| error::UnexpectedSnafu { violated: "payload budget closed" }.build()), + } + } + + /// Accounts for retained input allocations, conversion capacity growth and + /// array metadata. Pass zero backing when the caller detaches oversized slices. + pub(crate) fn conversion_budget( + backing: usize, + columns: usize, + requested: usize, + ) -> Result<(usize, usize)> { + let overhead = columns.saturating_mul(1024); + let available = PAYLOAD_BYTES + .saturating_sub(backing) + .saturating_sub(overhead) + / 4; + let conversion = requested.min(available); + ensure!( + conversion > 0, + error::LogicalTableExportResourceSnafu { + reason: "batch backing buffers exceed request payload budget" + } + ); + Ok((conversion, backing.saturating_add(overhead))) + } +} + +pub(crate) struct Payload { + pub(crate) batch: RecordBatch, + pub(crate) permit: OwnedSemaphorePermit, +} + +pub(crate) struct TableWriters { + pub(crate) packed: Option, + current: Option<(u32, mpsc::Sender)>, + tasks: FuturesUnordered>>, + budget: Arc, +} + +impl TableWriters { + pub(crate) fn new(budget: Arc) -> Self { + Self { + current: None, + packed: None, + tasks: FuturesUnordered::new(), + budget, + } + } + + #[cfg(test)] + pub(crate) fn pending_tasks(&self) -> usize { + self.tasks.len() + } + + pub(crate) fn table_id(&self) -> Option { + self.current.as_ref().map(|(id, _)| *id) + } + + pub(crate) fn close_input(&mut self) { + self.current = None; + } + + pub(crate) async fn open( + &mut self, + id: u32, + table: &LogicalTableProjection, + store: &ObjectStore, + limits: LogicalTableExportLimits, + token: &CancellationToken, + ) -> Result<()> { + // EOF must precede acquiring the next slot, including when P is one. + self.close_input(); + let permit = self.budget.writer(token).await?; + self.reap_for_admission(token).await?; + let writer = if let Some(packed) = &self.packed { + ActiveWriter::open_packed(table, id, store, limits, packed.clone())? + } else { + ActiveWriter::open(table, store, limits).await? + }; + let (sender, receiver) = mpsc::channel(2); + self.current = Some((id, sender)); + let token = token.clone(); + self.tasks.push(common_runtime::spawn_global(async move { + let _permit = permit; + let guard = token.clone().drop_guard(); + let result = run_writer(writer, receiver, &token).await; + if result.is_ok() { + guard.disarm(); + } + result + })); + Ok(()) + } + + async fn reap_for_admission(&mut self, token: &CancellationToken) -> Result<()> { + while let Some(Some(result)) = self.tasks.next().now_or_never() { + result.context(error::JoinTaskSnafu)??; + } + // A worker can release its permit before its JoinHandle becomes ready. + while self.tasks.len() > self.budget.max_writers { + let result = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + result = self.tasks.next() => result.context(error::UnexpectedSnafu { violated: "writer task queue unexpectedly empty" })?, + }; + result.context(error::JoinTaskSnafu)??; + } + Ok(()) + } + + pub(crate) async fn send(&self, payload: Payload, token: &CancellationToken) -> Result<()> { + let (_, sender) = self.current.as_ref().context(error::UnexpectedSnafu { + violated: "missing logical writer", + })?; + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + result = sender.send(payload) => result.map_err(|_| error::LogicalTableExportCancelledSnafu.build()), + } + } + + pub(crate) async fn drain( + &mut self, + result: Result<()>, + token: &CancellationToken, + ) -> Result<()> { + self.close_input(); + let mut first_error = result.err(); + if first_error.is_some() { + token.cancel(); + } + while let Some(result) = self.tasks.next().await { + if let Err(err) = result.context(error::JoinTaskSnafu).and_then(|r| r) { + retain_error(&mut first_error, err); + token.cancel(); + } + } + match first_error { + Some(err) => Err(err), + None => Ok(()), + } + } +} + +pub(crate) fn retain_error(first: &mut Option, error: error::Error) { + if first.as_ref().is_none_or(|err| { + matches!( + err, + error::Error::LogicalTableExportCancelled { .. } + | error::Error::DatabaseExportCancelled { .. } + ) + }) { + *first = Some(error); + } +} + +async fn run_writer( + mut writer: ActiveWriter, + mut receiver: mpsc::Receiver, + token: &CancellationToken, +) -> Result<()> { + let result = async { + loop { + let payload = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + payload = receiver.recv() => payload, + }; + let Some(Payload { batch, permit }) = payload else { + break; + }; + let result = writer.writer.write(batch, Some(token)).await; + drop(permit); + result.map_err(|error| map_writer_error(error, &writer.path))?; + } + writer + .writer + .finish(Some(token)) + .await + .map_err(|error| map_writer_error(error, &writer.path))?; + check_cancelled(token) + } + .await; + if result.is_err() { + token.cancel(); + receiver.close(); + drop(receiver); + if let Err(error) = writer.writer.abort().await { + common_telemetry::warn!(error; "Failed to abort Metric export file"); + } + } + result +} + +#[cfg(test)] +mod tests { + use tokio::sync::oneshot; + + use super::*; + + #[tokio::test] + async fn released_permit_does_not_hide_pending_join_handles() { + for parallelism in [1, 4] { + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut writers = TableWriters::new(budget.clone()); + let mut resume = Vec::new(); + for _ in 0..=parallelism { + let permit = budget.writer(&token).await.unwrap(); + let (released_tx, released_rx) = oneshot::channel(); + let (resume_tx, resume_rx) = oneshot::channel(); + writers.tasks.push(common_runtime::spawn_global(async move { + drop(permit); + released_tx.send(()).unwrap(); + resume_rx.await.unwrap(); + Ok(()) + })); + released_rx.await.unwrap(); + resume.push(resume_tx); + } + assert_eq!(budget.available().0, parallelism); + + let mut reap = Box::pin(writers.reap_for_admission(&token)); + assert!(reap.as_mut().now_or_never().is_none()); + resume.pop().unwrap().send(()).unwrap(); + reap.await.unwrap(); + assert_eq!(writers.pending_tasks(), parallelism); + + for tx in resume { + tx.send(()).unwrap(); + } + writers.drain(Ok(()), &token).await.unwrap(); + assert_eq!(budget.available().0, parallelism); + } + } +} diff --git a/src/operator/src/test_util/partition_manager.rs b/src/operator/src/test_util/partition_manager.rs index e51ac49051e..f80e2cc4fe2 100644 --- a/src/operator/src/test_util/partition_manager.rs +++ b/src/operator/src/test_util/partition_manager.rs @@ -368,3 +368,127 @@ async fn test_logical_split_rows_rejects_physical_only_partition_column() { partition::error::Error::UndefinedColumn { ref column, .. } if column == "b" )); } + +#[tokio::test] +async fn test_partition_routing_with_empty_default_expression() { + use datatypes::arrow::array::{ArrayRef, RecordBatch, StringArray}; + use datatypes::value::Value; + use partition::expr::col; + use store_api::storage::RegionId; + + for (transition, target_count) in [(true, 1), (true, 3), (false, 3)] { + let backend = Arc::new(common_meta::kv_backend::memory::MemoryKvBackend::new()); + let route_cache = test_new_table_route_cache(backend.clone()); + let manager = PartitionRuleManager::new( + backend.clone(), + route_cache.clone(), + test_new_partition_info_cache(route_cache), + ); + let value = |s: &str| Value::String(s.into()); + let mut exprs = vec![ + col("host").lt(value("36z0")), + col("host") + .gt_eq(value("36z0")) + .and(col("host").lt(value("7Ez0"))), + col("host") + .gt_eq(value("7Ez0")) + .and(col("host").lt(value("BMz0"))), + col("host").gt_eq(value("BMz0")), + ]; + if target_count == 1 { + exprs = vec![ + col("host").lt(value("36z0")), + col("host").gt_eq(value("36z0")), + ]; + } + let routes = exprs + .into_iter() + .enumerate() + .map(|(i, expr)| RegionRoute { + region: Region { + id: RegionId::new(1027, i as u32), + partition_expr: if transition && i == 0 { + String::new() + } else { + expr.as_json_str().unwrap() + }, + ..Default::default() + }, + leader_peer: Some(Peer::new(1, "")), + ..Default::default() + }) + .collect(); + let info = new_test_table_info_with_columns( + 1027, + "physical", + vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ], + vec![0], + ); + let metadata_manager = TableMetadataManager::new(backend); + metadata_manager + .create_table_metadata( + info.clone(), + TableRouteValue::physical(routes), + new_test_region_wal_options((0..=target_count).collect()), + ) + .await + .unwrap(); + let mut logical_info = info.clone(); + logical_info.ident.table_id = 1028; + logical_info.name = "logical".into(); + metadata_manager + .create_table_metadata( + logical_info.clone(), + TableRouteValue::logical(1027), + HashMap::new(), + ) + .await + .unwrap(); + + let hosts = [None, Some("1"), Some("5"), Some("8"), Some("animi")]; + let expected_regions = if target_count == 1 { + [0, 0, 1, 1, 1] + } else { + [0, 0, 1, 2, 3] + }; + let batch = RecordBatch::try_from_iter([( + "host", + Arc::new(StringArray::from(hosts.to_vec())) as ArrayRef, + )]) + .unwrap(); + for table_info in [&info, &logical_info] { + let (rule, _) = manager.find_table_partition_rule(table_info).await.unwrap(); + for (host, expected) in hosts.into_iter().zip(expected_regions) { + assert_eq!( + rule.find_region(&[host.map(value).unwrap_or(Value::Null)]) + .unwrap(), + expected + ); + } + let splits = rule.split_record_batch(&batch).unwrap(); + assert_eq!(splits.len(), target_count as usize + 1); + for (region, mask) in splits { + let selected = mask + .array() + .iter() + .enumerate() + .filter_map(|(i, selected)| (selected == Some(true)).then_some(i)) + .collect::>(); + let expected = expected_regions + .iter() + .enumerate() + .filter_map(|(i, expected)| (*expected == region).then_some(i)) + .collect::>(); + assert_eq!(selected, expected); + } + } + } +} diff --git a/src/partition/Cargo.toml b/src/partition/Cargo.toml index a8e3a8ae11a..6773cc5878b 100644 --- a/src/partition/Cargo.toml +++ b/src/partition/Cargo.toml @@ -9,18 +9,14 @@ workspace = true [dependencies] api.workspace = true -async-trait.workspace = true common-base.workspace = true common-error.workspace = true common-macro.workspace = true common-meta.workspace = true -common-query.workspace = true -common-telemetry.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true datafusion-physical-expr.workspace = true datatypes.workspace = true -itertools.workspace = true moka = { workspace = true, features = ["future"] } serde.workspace = true serde_json.workspace = true diff --git a/src/partition/src/manager.rs b/src/partition/src/manager.rs index 89bb61c4174..68c9586e908 100644 --- a/src/partition/src/manager.rs +++ b/src/partition/src/manager.rs @@ -206,17 +206,25 @@ impl PartitionRuleManager { .iter() .map(|r| (r.id.region_number(), r.partition_expr_version)) .collect::>>(); - let regions = partition_info + let (mut regions, exprs): (Vec<_>, Vec<_>) = partition_info .partitions .iter() - .map(|x| x.id.region_number()) - .collect::>(); - let exprs = partition_info - .partitions - .iter() - .filter_map(|x| x.partition_expr.as_ref()) - .cloned() - .collect::>(); + .filter_map(|partition| { + partition + .partition_expr + .as_ref() + .map(|expr| (partition.id.region_number(), expr.clone())) + }) + .unzip(); + // Empty expressions belong to the default region. Keep it after the + // explicit regions so expression indices still identify the right region. + regions.extend( + partition_info + .partitions + .iter() + .filter(|partition| partition.partition_expr.is_none()) + .map(|partition| partition.id.region_number()), + ); let partition_rule = Arc::new(MultiDimPartitionRule::try_new( partition_columns, regions, diff --git a/src/pipeline/Cargo.toml b/src/pipeline/Cargo.toml index 620ac36ac77..9ebe51ecf24 100644 --- a/src/pipeline/Cargo.toml +++ b/src/pipeline/Cargo.toml @@ -14,28 +14,22 @@ ahash.workspace = true api.workspace = true arrow.workspace = true arrow-schema.workspace = true -async-trait.workspace = true catalog.workspace = true chrono.workspace = true chrono-tz.workspace = true common-catalog.workspace = true common-error.workspace = true -common-function.workspace = true +common-frontend.workspace = true common-macro.workspace = true common-meta.workspace = true common-query.workspace = true common-recordbatch.workspace = true -common-runtime.workspace = true common-telemetry.workspace = true common-time.workspace = true -crossbeam-utils.workspace = true csv = "1.3" -dashmap.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true -datafusion-functions.workspace = true -datafusion-physical-expr.workspace = true datatypes.workspace = true dyn-fmt = "0.4" enum_dispatch = "0.3" @@ -50,7 +44,6 @@ moka = { workspace = true, features = ["future"] } once_cell.workspace = true operator.workspace = true ordered-float.workspace = true -paste.workspace = true prometheus.workspace = true query.workspace = true regex.workspace = true @@ -68,8 +61,6 @@ yaml-rust = "0.4" [dev-dependencies] catalog = { workspace = true, features = ["testing"] } criterion = { workspace = true, features = ["html_reports"] } -rayon = "1.0" -session = { workspace = true, features = ["testing"] } [[bench]] name = "processor" diff --git a/src/pipeline/src/etl.rs b/src/pipeline/src/etl.rs index f83b5c2626c..f78a515fb59 100644 --- a/src/pipeline/src/etl.rs +++ b/src/pipeline/src/etl.rs @@ -278,24 +278,6 @@ impl PipelineExecOutput { } } - // Backward compatibility helper that returns first ContextOpt with all its rows - // or merges all rows with default ContextOpt for multi-context scenarios - pub fn into_legacy_format(self) -> Option<(ContextOpt, Vec)> { - if let Self::Transformed(TransformedOutput { rows_by_context }) = self { - if rows_by_context.len() == 1 { - let (opt, rows) = rows_by_context.into_iter().next().unwrap(); - Some((opt, rows)) - } else { - // Multiple contexts: merge all rows with default ContextOpt for test compatibility - let all_rows: Vec = - rows_by_context.into_values().flatten().collect(); - Some((ContextOpt::default(), all_rows)) - } - } else { - None - } - } - // Note: This is a test only function, do not use it in production. pub fn into_dispatched(self) -> Option { if let Self::DispatchedTo(d, _) = self { diff --git a/src/pipeline/src/etl/transform/index.rs b/src/pipeline/src/etl/transform/index.rs index 53e16dca222..b0cdd9c4fe0 100644 --- a/src/pipeline/src/etl/transform/index.rs +++ b/src/pipeline/src/etl/transform/index.rs @@ -25,7 +25,7 @@ const INDEX_INVERTED: &str = "inverted"; #[allow(clippy::enum_variant_names)] pub enum Index { Time, - // deprecated, use Inverted instead + // Deprecated, equivalent to `tag: true`. Tag, Fulltext, Skipping, diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index bbe3949a540..16c9e3399b6 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -741,9 +741,9 @@ fn resolve_value( &ConcreteDataType::string_datatype(), schema_info, )?; - Some(ValueData::StringValue(String::from_utf8_lossy_owned( - v.to_vec(), - ))) + Some(ValueData::StringValue( + String::from_utf8_lossy(&v).into_owned(), + )) } VrlValue::Regex(v) => { @@ -911,7 +911,7 @@ fn identity_pipeline_inner( } let column_count = schema_info.schema.len(); - for (_, row) in opt_map.iter_mut() { + for row in opt_map.values_mut() { for row in row.iter_mut() { assert!( column_count >= row.values.len(), diff --git a/src/pipeline/src/lib.rs b/src/pipeline/src/lib.rs index 5d1acb7462e..15954b18f98 100644 --- a/src/pipeline/src/lib.rs +++ b/src/pipeline/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(string_from_utf8_lossy_owned)] - mod dispatcher; pub mod error; mod etl; diff --git a/src/pipeline/src/manager/pipeline_cache.rs b/src/pipeline/src/manager/pipeline_cache.rs index e60963e7f84..f0e9ca52970 100644 --- a/src/pipeline/src/manager/pipeline_cache.rs +++ b/src/pipeline/src/manager/pipeline_cache.rs @@ -16,6 +16,7 @@ use std::future::Future; use std::sync::Arc; use std::time::Duration; +use common_frontend::metrics; use datatypes::timestamp::TimestampNanosecond; use moka::future::Cache; @@ -32,7 +33,7 @@ const PIPELINES_CACHE_SIZE: u64 = 10000; /// to encapsulate inner cache. Only public methods are exposed. /// /// `pipelines` and `original_pipelines` are keyed by the *requested* schema so -/// a lookup is a single key probe, as [`Cache::try_get_with`] requires; +/// a lookup is a single key probe through [`Cache::entry`]; /// resolving it to a stored schema is the loader's job. `failover_cache` has no /// loader and keeps the stored-schema key. pub(crate) struct PipelineCache { @@ -71,7 +72,7 @@ impl PipelineCache { } } - /// Concurrent misses on the same key share one `init` call. + /// Concurrent misses share one `init` call; successful waiters count as hits. pub(crate) async fn get_pipeline_with( &self, schema: &str, @@ -80,13 +81,17 @@ impl PipelineCache { init: impl Future>>, ) -> Result> { let key = generate_pipeline_cache_key(schema, name, version); - self.pipelines - .try_get_with(key, init) - .await + let entry = self.pipelines.entry(key).or_try_insert_with(init).await; + metrics::record_cache_lookup( + "pipeline", + entry.as_ref().is_ok_and(|entry| !entry.is_fresh()), + ); + entry + .map(|entry| entry.into_value()) .map_err(|error| CacheLoadSnafu { error }.build()) } - /// Concurrent misses on the same key share one `init` call. + /// Concurrent misses share one `init` call; successful waiters count as hits. pub(crate) async fn get_pipeline_str_with( &self, schema: &str, @@ -95,9 +100,17 @@ impl PipelineCache { init: impl Future>, ) -> Result { let key = generate_pipeline_cache_key(schema, name, version); - self.original_pipelines - .try_get_with(key, init) - .await + let entry = self + .original_pipelines + .entry(key) + .or_try_insert_with(init) + .await; + metrics::record_cache_lookup( + "pipeline_original", + entry.as_ref().is_ok_and(|entry| !entry.is_fresh()), + ); + entry + .map(|entry| entry.into_value()) .map_err(|error| CacheLoadSnafu { error }.build()) } @@ -114,6 +127,7 @@ impl PipelineCache { generate_pipeline_cache_key(schema, name, version), ] { if let Some(content) = self.failover_cache.get(&key).await { + metrics::record_cache_lookup("pipeline_failover", true); return Ok(Some(content)); } } @@ -126,7 +140,7 @@ impl PipelineCache { .filter(|(k, _)| k.ends_with(&suffix)) .collect::>(); - match found.len() { + let result = match found.len() { 0 => Ok(None), 1 => Ok(Some(found.remove(0).1)), _ => MultiPipelineWithDiffSchemaSnafu { @@ -139,7 +153,12 @@ impl PipelineCache { .join(","), } .fail(), - } + }; + metrics::record_cache_lookup( + "pipeline_failover", + result.as_ref().is_ok_and(|content| content.is_some()), + ); + result } pub(crate) async fn insert_failover_cache(&self, content: PipelineContent, with_latest: bool) { @@ -191,7 +210,8 @@ impl PipelineCache { mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; - use tokio::sync::Barrier; + use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; + use tokio::sync::Notify; use super::*; @@ -205,36 +225,149 @@ mod tests { } } - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[tokio::test] async fn test_concurrent_misses_run_one_loader() { const CONCURRENCY: usize = 8; - let cache = Arc::new(PipelineCache::new(Duration::from_secs(60))); - let loads = Arc::new(AtomicUsize::new(0)); - let barrier = Arc::new(Barrier::new(CONCURRENCY)); - - let handles = (0..CONCURRENCY) + let hits = CACHE_HIT.with_label_values(&["pipeline_original"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_original"]); + let before = (hits.get(), misses.get()); + let cache = PipelineCache::new(Duration::from_secs(60)); + let loads = AtomicUsize::new(0); + let release = Notify::new(); + let mut requests = (0..CONCURRENCY) .map(|_| { - let (cache, loads, barrier) = (cache.clone(), loads.clone(), barrier.clone()); - tokio::spawn(async move { - barrier.wait().await; + Box::pin(cache.get_pipeline_str_with("db", "p", None, async { + loads.fetch_add(1, Ordering::SeqCst); + release.notified().await; + Ok(content_at(1)) + })) + }) + .collect::>(); + for request in &mut requests { + assert!(futures::poll!(request.as_mut()).is_pending()); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + release.notify_waiters(); + for request in requests { + assert_eq!(request.await.unwrap(), content_at(1)); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!( + (hits.get(), misses.get()), + (before.0 + CONCURRENCY as u64 - 1, before.1 + 1) + ); + } + + #[tokio::test] + async fn test_pipeline_lookup_metrics() { + let cache = PipelineCache::new(Duration::from_secs(60)); + for compiled in [false, true] { + let cache_type = if compiled { + "pipeline" + } else { + "pipeline_original" + }; + let hits = CACHE_HIT.with_label_values(&[cache_type]); + let misses = CACHE_MISS.with_label_values(&[cache_type]); + let before = (hits.get(), misses.get()); + let load = async |fail| { + if compiled { + cache + .get_pipeline_with("db", "p", None, async { + if fail { + return Err(crate::error::PipelineNotFoundSnafu { + name: "p", + version: None, + } + .build()); + } + Ok(Arc::new( + crate::manager::table::PipelineTable::compile_pipeline( + "transform:\n - fields: [message]\n type: string", + )?, + )) + }) + .await + .map(|_| ()) + } else { cache .get_pipeline_str_with("db", "p", None, async { - loads.fetch_add(1, Ordering::SeqCst); - // Hold the loader open so every caller is waiting on it. - tokio::time::sleep(Duration::from_millis(100)).await; + if fail { + return Err(crate::error::PipelineNotFoundSnafu { + name: "p", + version: None, + } + .build()); + } Ok(content_at(1)) }) .await - .unwrap() - }) - }) - .collect::>(); - - for handle in handles { - assert_eq!(handle.await.unwrap(), content_at(1)); + .map(|_| ()) + } + }; + assert!(load(true).await.is_err()); + load(false).await.unwrap(); + // A warm lookup must not evaluate the failing initializer. + load(true).await.unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2)); } - assert_eq!(loads.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_failover_lookup_metrics() { + let hits = CACHE_HIT.with_label_values(&["pipeline_failover"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_failover"]); + let before = (hits.get(), misses.get()); + let cache = PipelineCache::new(Duration::from_secs(60)); + assert!( + cache + .get_failover_cache("db", "p", None) + .await + .unwrap() + .is_none() + ); + let local = PipelineContent { + schema: "a".into(), + ..content_at(1) + }; + cache.insert_failover_cache(local.clone(), true).await; + // Direct schema lookup and cross-schema fallback each count one hit. + assert_eq!( + cache.get_failover_cache("a", "p", None).await.unwrap(), + Some(local.clone()) + ); + assert_eq!( + cache.get_failover_cache("db", "p", None).await.unwrap(), + Some(local) + ); + cache + .insert_failover_cache( + PipelineContent { + schema: "b".into(), + ..content_at(2) + }, + true, + ) + .await; + assert!(cache.get_failover_cache("db", "p", None).await.is_err()); + // Global pipelines win over local entries inserted before or after them. + for (schema, version) in [(EMPTY_SCHEMA_NAME, 3), ("c", 4)] { + cache + .insert_failover_cache( + PipelineContent { + schema: schema.into(), + ..content_at(version) + }, + true, + ) + .await; + assert_eq!( + cache.get_failover_cache("db", "p", None).await.unwrap(), + Some(content_at(3)) + ); + } + assert_eq!((hits.get(), misses.get()), (before.0 + 4, before.1 + 2)); } #[tokio::test] @@ -288,25 +421,4 @@ mod tests { let failover = cache.get_failover_cache("b", "p", None).await.unwrap(); assert_eq!(failover.map(|c| c.version), Some(v2.version)); } - - #[tokio::test] - async fn test_failover_serves_global_pipeline_to_unwarmed_schema() { - let cache = PipelineCache::new(Duration::from_secs(60)); - let content = content_at(1); - - cache.insert_failover_cache(content.clone(), true).await; - - let found = cache.get_failover_cache("b", "p", None).await.unwrap(); - assert_eq!(found, Some(content.clone())); - - // A same-named pipeline under another schema must not shadow the global one. - let schema_local = PipelineContent { - schema: "x".to_string(), - ..content_at(2) - }; - cache.insert_failover_cache(schema_local, true).await; - - let found = cache.get_failover_cache("b", "p", None).await.unwrap(); - assert_eq!(found, Some(content)); - } } diff --git a/src/pipeline/src/manager/pipeline_operator.rs b/src/pipeline/src/manager/pipeline_operator.rs index c9731b36f20..cace74171c2 100644 --- a/src/pipeline/src/manager/pipeline_operator.rs +++ b/src/pipeline/src/manager/pipeline_operator.rs @@ -19,6 +19,7 @@ use std::time::{Duration, Instant}; use api::v1::CreateTableExpr; use catalog::{CatalogManagerRef, RegisterSystemTableRequest}; use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, default_engine}; +use common_frontend::metrics; use common_meta::rpc::ddl::TriggerReason; use common_telemetry::info; use common_time::FOREVER; @@ -164,7 +165,9 @@ impl PipelineOperator { /// Get a pipeline table from the cache. pub fn get_pipeline_table_from_cache(&self, catalog: &str) -> Option { - self.tables.read().unwrap().get(catalog).cloned() + let table = self.tables.read().unwrap().get(catalog).cloned(); + metrics::record_cache_lookup("pipeline_table", table.is_some()); + table } } diff --git a/src/pipeline/src/tablesuffix.rs b/src/pipeline/src/tablesuffix.rs index 01ce4202cfc..4a72f49bbc6 100644 --- a/src/pipeline/src/tablesuffix.rs +++ b/src/pipeline/src/tablesuffix.rs @@ -56,7 +56,7 @@ impl TableSuffixTemplate { let v = val.get(key.as_str())?; match v { VrlValue::Integer(v) => Some(v.to_string()), - VrlValue::Bytes(v) => Some(String::from_utf8_lossy_owned(v.to_vec())), + VrlValue::Bytes(v) => Some(String::from_utf8_lossy(v).into_owned()), _ => None, } }) diff --git a/src/promql/benches/bench_range_fn.rs b/src/promql/benches/bench_range_fn.rs index 0f0a4867665..51e0aefdc79 100644 --- a/src/promql/benches/bench_range_fn.rs +++ b/src/promql/benches/bench_range_fn.rs @@ -17,25 +17,14 @@ use std::sync::Arc; use criterion::{BenchmarkId, Criterion, criterion_group}; -use datafusion::arrow::array::{ - Array, ArrayRef, DictionaryArray, Float64Array, TimestampMillisecondArray, -}; -use datafusion::arrow::datatypes::{ - DataType as ArrowDataType, Field as ArrowField, Int64Type, Schema, -}; -use datafusion::arrow::record_batch::RecordBatch; -use datafusion::common::ToDFSchema; -use datafusion::datasource::memory::MemorySourceConfig; -use datafusion::datasource::source::DataSourceExec; -use datafusion::execution::context::TaskContext; -use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; -use datafusion::physical_plan::{ColumnarValue, ExecutionPlan}; +use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray}; +use datafusion::arrow::datatypes::TimeUnit; +use datafusion::physical_plan::ColumnarValue; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::ScalarFunctionArgs; use datatypes::arrow::datatypes::{DataType, Field}; -use futures::StreamExt; -use promql::extension_plan::RangeManipulate; +use promql::extension_plan::RangeManipulateStream; use promql::functions::{ AbsentOverTime, Changes, CountOverTime, Delta, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, @@ -270,7 +259,7 @@ fn make_quantile_input(num_points: usize, window_size: u32) -> Vec Vec { - let (ts_range, val_range, _) = build_sliding_ranges( + let (ts_range, val_range, eval_ts) = build_sliding_ranges( num_points, window_size, 1, @@ -282,6 +271,8 @@ fn make_predict_linear_input(num_points: usize, window_size: u32) -> Vec, + window_ms: i64, + interval_ms: i64, + start_ms: i64, + end_ms: i64, +} -fn make_range_manipulate_batch(timestamps: Vec, field_count: usize) -> RecordBatch { - let mut fields = Vec::with_capacity(field_count + 1); - let mut columns: Vec = Vec::with_capacity(field_count + 1); - fields.push(ArrowField::new( - "timestamp", - ArrowDataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), - false, - )); - columns.push(Arc::new(TimestampMillisecondArray::from(timestamps)) as _); - - for field_index in 0..field_count { - fields.push(ArrowField::new( - format!("value_{field_index}"), - ArrowDataType::Float64, - false, - )); - let values = (0..columns[0].len()) - .map(|row| row as f64 + field_index as f64 * 0.01) - .collect::>(); - columns.push(Arc::new(Float64Array::from(values)) as _); +impl CursorScanWorkload { + /// The workload through the narrow `i64` carrier, or `None` when production would fall + /// back to `i128` for it. + /// + /// For the fallback cases the fast path is timed on a twin workload: the same shape + /// translated onto a timeline the narrow carrier holds, which keeps the number of sample + /// and window comparisons identical to the fallback run. + fn bench_scan_narrow(&self) -> Option> { + RangeManipulateStream::bench_scan_narrow( + self.time_unit, + self.offset_ms, + self.window_ms, + self.interval_ms, + &self.timestamps, + self.start_ms, + self.end_ms, + ) } - RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() -} + /// The workload through the exact `i128` carrier, the fallback production takes + /// whenever [`Self::bench_scan_narrow`] returns `None`. + fn bench_scan_wide(&self) -> Vec<(u32, u32)> { + RangeManipulateStream::bench_scan_wide( + self.time_unit, + self.offset_ms, + self.window_ms, + self.interval_ms, + &self.timestamps, + self.start_ms, + self.end_ms, + ) + } -async fn execute_and_drain( - plan: Arc, - task_ctx: Arc, -) -> Vec { - let mut stream = plan.execute(0, task_ctx).unwrap(); - let mut output = Vec::new(); - while let Some(batch) = stream.next().await { - output.push(batch.unwrap()); + fn evaluations(&self) -> usize { + ((self.end_ms - self.start_ms) / self.interval_ms) as usize + 1 + } + + fn label(&self) -> String { + format!( + "N{}_eval{}_window{}x{}ms", + self.timestamps.len(), + self.evaluations(), + self.window_ms / self.interval_ms, + self.interval_ms + ) } - output } -fn range_keys(batch: &RecordBatch, index: usize) -> Vec<(u32, u32)> { - let ranges = batch - .column(index) - .as_any() - .downcast_ref::>() - .unwrap() - .clone(); - RangeArray::try_new(ranges) - .unwrap() - .ranges() - .map(Option::unwrap) - .collect() +/// Cadence of the cursor-scan benchmarks: PromQL's usual 15s step. +const CURSOR_SCAN_CADENCE_MS: i64 = 15_000; + +/// Milliseconds of a one-hour window, the widest window benchmarked here. +const CURSOR_SCAN_HOUR_MS: i64 = 3_600_000; + +/// Builds a workload from a first timestamp and a sample count. +/// +/// `sample_spacing_native` is the distance between samples in native ticks: 15s worth, so +/// every native-precision workload has the same shape (one sample per cadence, 20 samples +/// per 300s window). The aligned instants start at the first shifted sample, so every +/// window overlaps the input the same way in both arms of the A/B comparison. +fn cursor_scan_workload( + time_unit: TimeUnit, + offset_ms: i64, + first_timestamp: i64, + samples: usize, + window_ms: i64, + sample_spacing_native: i64, +) -> CursorScanWorkload { + let interval_ms = CURSOR_SCAN_CADENCE_MS; + let timestamps = (0..samples as i64) + .map(|index| first_timestamp + index * sample_spacing_native) + .collect::>(); + // `Millisecond` instants of the first shifted sample, flooring like `calculate_range`. + let start_ms = shifted_floor_millis(time_unit, offset_ms, first_timestamp); + let evaluations = samples - 1; + CursorScanWorkload { + time_unit, + offset_ms, + timestamps, + window_ms, + interval_ms, + start_ms, + end_ms: start_ms + evaluations as i64 * interval_ms, + } } -fn brute_force_range_keys( - timestamps: &[i64], - evaluations: usize, - window_points: usize, -) -> Vec<(u32, u32)> { - let range = window_points as i64 * RANGE_MANIPULATE_CADENCE_MS; - (0..evaluations) - .map(|evaluation| { - let current = evaluation as i64 * RANGE_MANIPULATE_CADENCE_MS; - let mut offset = None; - let mut length = 0; - for (index, timestamp) in timestamps.iter().enumerate() { - if *timestamp > current - range && *timestamp <= current { - offset.get_or_insert(index as u32); - length += 1; - } - } - (offset.unwrap_or(0), length) - }) - .collect() +/// Native ticks of one benchmark cadence, i.e. 15s in the given precision. +fn cadence_in_native_ticks(time_unit: TimeUnit) -> i64 { + match time_unit { + TimeUnit::Second => CURSOR_SCAN_CADENCE_MS / 1_000, + TimeUnit::Millisecond => CURSOR_SCAN_CADENCE_MS, + TimeUnit::Microsecond => CURSOR_SCAN_CADENCE_MS * 1_000, + TimeUnit::Nanosecond => CURSOR_SCAN_CADENCE_MS * 1_000_000, + } } -fn validate_range_manipulate_output( - output: &[RecordBatch], - field_count: usize, - expected_range_keys: Option<&[(u32, u32)]>, -) { - let Some(expected_range_keys) = expected_range_keys else { - assert!(output.is_empty()); - return; +/// Floor of a shifted sample in milliseconds, in exact `i128` native precision. +fn shifted_floor_millis(time_unit: TimeUnit, offset_ms: i64, timestamp: i64) -> i64 { + let scale = match time_unit { + TimeUnit::Second => 1_000_000_000i128, + TimeUnit::Millisecond => 1_000_000i128, + TimeUnit::Microsecond => 1_000i128, + TimeUnit::Nanosecond => 1i128, + }; + i64::try_from( + ((timestamp as i128) * scale + (offset_ms as i128) * 1_000_000).div_euclid(1_000_000), + ) + .expect("the shifted millisecond must stay legal") +} + +/// A/B benchmark of the cursor-scan carriers behind `RangeManipulate::scan_ranges`. +/// +/// `RangeManipulate` scans the aligned instants in `i64` nanoseconds whenever +/// `narrow_scan` proves the batch and the bounds fit, and in exact `i128` nanoseconds +/// otherwise. Both arms below drive the production scan through the carrier the operator +/// would pick, so the numbers isolate the carrier and not the surrounding operator work: +/// +/// - `ms_zero_offset`: millisecond batch with no offset, the `i64` fast path. The same +/// workload is also timed through `i128` to price the fallback on this input. +/// - `ms_zero_offset_long_window`: same, with a one-hour window, where the scan advances +/// both cursors much further per evaluation. +/// - `ms_large_offset`: an offset whose nanosecond product overflows `i64`, so production +/// falls back to `i128`. The `i64` arm times the identical workload shape translated onto +/// a narrow timeline, which keeps the selections and thus the work identical. +/// - `ns_native_floor`: nanosecond batch at the `i64::MIN` floor, where the shifted start +/// bound underflows `i64` nanoseconds, so production falls back to `i128`. The `i64` arm +/// again times the translated narrow shape. +/// +/// Every arm asserts that both carriers selected the same ranges, so the comparison cannot +/// be won by doing less work. The fallback triggers themselves are pinned by the unit tests +/// in `range_manipulate.rs` (`scan_falls_back_to_i128_when_only_the_sample_tail_overflows` +/// and `scan_carriers_split_at_the_nanosecond_boundaries`). +fn bench_cursor_scan_carriers(c: &mut Criterion) { + const SAMPLES: usize = 4_096; + const SHORT_WINDOW_MS: i64 = 20 * CURSOR_SCAN_CADENCE_MS; + // `i64::MAX` milliseconds is ~9.22e18 ms; one microsecond past it in the nanosecond + // product makes `narrow_scan` reject the offset, while the shifted millisecond stays legal. + const LARGE_OFFSET_MS: i64 = 9_223_372_036_855; + const LARGE_OFFSET_BASE_MS: i64 = 1_000_000_000_000_000; + const NANOSECOND_FLOOR_SHIFT_NS: i64 = 9_223_372_036_000_000_000; + const NANOSECOND_FLOOR_SAMPLES: usize = 4_096; + + let ms_zero_offset = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + let ms_zero_offset_long_window = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + CURSOR_SCAN_HOUR_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + // The large offset is shifted inside one millisecond of the nanosecond ceiling, so the + // payload stays legal: the fallback is driven by the product, not by an illegal instant. + let ms_large_offset = cursor_scan_workload( + TimeUnit::Millisecond, + LARGE_OFFSET_MS, + LARGE_OFFSET_BASE_MS, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + // The narrow twin of `ms_large_offset`: the same batch shape without the offset. + let ms_large_offset_narrow = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + + // Nanosecond batch on the native floor: samples at `i64::MIN + i * 15s`, whose shifted + // start bound underflows `i64` nanoseconds. + let nanosecond_floor_first = i64::MIN; + let nanosecond_floor = cursor_scan_workload( + TimeUnit::Nanosecond, + 0, + nanosecond_floor_first, + NANOSECOND_FLOOR_SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Nanosecond), + ); + // The narrow twin: the same shape translated by a multiple of one millisecond, so both + // the samples and the aligned instants keep their relative geometry. + let nanosecond_floor_narrow = CursorScanWorkload { + time_unit: TimeUnit::Nanosecond, + offset_ms: 0, + timestamps: nanosecond_floor + .timestamps + .iter() + .map(|timestamp| timestamp + NANOSECOND_FLOOR_SHIFT_NS) + .collect(), + window_ms: nanosecond_floor.window_ms, + interval_ms: nanosecond_floor.interval_ms, + start_ms: nanosecond_floor.start_ms + NANOSECOND_FLOOR_SHIFT_NS / 1_000_000, + end_ms: nanosecond_floor.end_ms + NANOSECOND_FLOOR_SHIFT_NS / 1_000_000, }; - assert_eq!(output.len(), 1); - let batch = &output[0]; - assert_eq!(batch.num_rows(), expected_range_keys.len()); - - let timestamp_range_keys = range_keys(batch, field_count + 1); - assert_eq!(timestamp_range_keys, expected_range_keys); - - for field_index in 0..field_count { - assert_eq!(range_keys(batch, field_index + 1), expected_range_keys); - } -} - -fn bench_range_manipulate_wall_time(c: &mut Criterion) { - let mut group = c.benchmark_group("range_manipulate_wall_time"); - let primary_points = 4_096; - let sparse_evaluations = primary_points - 1; - let mut cases = vec![ + let cases = [ ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 4, - 1, + "ms_zero_offset", + ms_zero_offset.clone(), + ms_zero_offset, false, ), ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 20, - 1, + "ms_zero_offset_long_window", + ms_zero_offset_long_window.clone(), + ms_zero_offset_long_window, false, ), ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 240, - 1, - false, - ), - ( - "regular_sparse", - (0..sparse_evaluations) - .step_by(2) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - sparse_evaluations, - 20, - 1, - false, - ), - ( - "all_empty_pathological", - (0..primary_points) - .map(|point| { - (primary_points as i64 + 21 + point as i64) * RANGE_MANIPULATE_CADENCE_MS - }) - .collect::>(), - primary_points, - 20, - 1, + "ms_large_offset", + ms_large_offset_narrow, + ms_large_offset, true, ), ( - "multi_field_control", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 20, - 4, - false, + "ns_native_floor", + nanosecond_floor_narrow, + nanosecond_floor, + true, ), ]; - for (case_name, timestamps, evaluations, window_points, field_count, expect_empty) in - cases.drain(..) - { - let expected_range_keys = (!expect_empty) - .then(|| brute_force_range_keys(×tamps, evaluations, window_points)); - let input_batch = make_range_manipulate_batch(timestamps, field_count); - let input_rows = input_batch.num_rows(); - let input_schema = input_batch.schema(); - let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: input_schema.clone().to_dfschema_ref().unwrap(), - }); - let field_columns = (0..field_count) - .map(|field_index| format!("value_{field_index}")) - .collect::>(); - let logical_plan = RangeManipulate::new( - 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, - logical_input, - ) - .unwrap(); - let physical_input: Arc = Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[vec![input_batch]], input_schema, None).unwrap(), - ))); - let execution_plan = logical_plan.to_execution_plan(physical_input); - let runtime = tokio::runtime::Runtime::new().unwrap(); - let task_ctx = datafusion::prelude::SessionContext::new().task_ctx(); + let mut group = c.benchmark_group("cursor_scan_carriers"); + group.warm_up_time(std::time::Duration::from_millis(500)); + group.measurement_time(std::time::Duration::from_secs(2)); + group.sample_size(50); - // Validate the public operator output and RangeArray keys before timing. - let output = runtime.block_on(execute_and_drain(execution_plan.clone(), task_ctx.clone())); - validate_range_manipulate_output(&output, field_count, expected_range_keys.as_deref()); + for (case_name, narrow, wide, wide_falls_back) in cases { + // Both arms must be timed on work the operator would really do: the narrow arm in + // `i64`, and the wide arm in `i128` exactly when the narrow baselines overflow. + let narrow_expected = narrow + .bench_scan_narrow() + .unwrap_or_else(|| panic!("{case_name}: the i64 arm must fit the narrow carrier")); + assert_eq!( + wide.bench_scan_narrow().is_none(), + wide_falls_back, + "{case_name}: the i128 arm's fallback trigger" + ); + // Fairness check before timing: both carriers have to select the same ranges, which + // makes the sample and window comparisons of the two arms identical. + assert_eq!( + wide.bench_scan_wide(), + narrow_expected, + "{case_name}: the two carriers must select the same ranges" + ); + assert_eq!( + narrow_expected.len(), + narrow.evaluations(), + "{case_name}: one range per aligned evaluation instant" + ); - // This measures public RangeManipulate execution wall time, including stream draining. + group.throughput(criterion::Throughput::Elements( + narrow.timestamps.len() as u64 + )); group.bench_with_input( - BenchmarkId::new( - case_name, - format!("N{input_rows}_eval{evaluations}_window{window_points}x15s"), - ), - &(), - |b, _| { - b.iter(|| { - let output = runtime - .block_on(execute_and_drain(execution_plan.clone(), task_ctx.clone())); - std::hint::black_box(output); - }) + BenchmarkId::new(format!("{case_name}/i64"), narrow.label()), + &narrow, + |b, workload| { + b.iter(|| std::hint::black_box(workload.bench_scan_narrow())); + }, + ); + group.bench_with_input( + BenchmarkId::new(format!("{case_name}/i128"), wide.label()), + &wide, + |b, workload| { + b.iter(|| std::hint::black_box(workload.bench_scan_wide())); }, ); } @@ -1194,5 +1256,5 @@ criterion_group!( bench_rate_window_steps, bench_edge_count_functions, bench_extrema_functions, - bench_range_manipulate_wall_time + bench_cursor_scan_carriers ); diff --git a/src/promql/src/extension_plan/absent.rs b/src/promql/src/extension_plan/absent.rs index fa811b3df7d..9afe85873c5 100644 --- a/src/promql/src/extension_plan/absent.rs +++ b/src/promql/src/extension_plan/absent.rs @@ -34,7 +34,7 @@ use datafusion::physical_plan::{ Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use datafusion_common::DFSchema; -use datafusion_expr::{EmptyRelation, col}; +use datafusion_expr::{EmptyRelation, ident}; use datatypes::arrow; use datatypes::arrow::array::{ArrayRef, Float64Array, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit}; @@ -108,7 +108,7 @@ impl UserDefinedLogicalNodeCore for Absent { return vec![]; } - vec![col(&self.time_index_column)] + vec![ident(&self.time_index_column)] } fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { diff --git a/src/promql/src/extension_plan/empty_metric.rs b/src/promql/src/extension_plan/empty_metric.rs index 91c8a1666a2..4d2e029417a 100644 --- a/src/promql/src/extension_plan/empty_metric.rs +++ b/src/promql/src/extension_plan/empty_metric.rs @@ -40,7 +40,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, StatisticsArgs, }; use datafusion::physical_planner::PhysicalPlanner; -use datafusion::prelude::{Expr, col, lit}; +use datafusion::prelude::{Expr, ident, lit}; use datafusion_expr::LogicalPlanBuilder; use datatypes::arrow::array::TimestampMillisecondArray; use datatypes::arrow::datatypes::SchemaRef; @@ -409,7 +409,7 @@ fn build_ts_only_schema(column_name: &str) -> DFSchema { pub fn build_special_time_expr(time_index_column_name: &str) -> Expr { let input_schema = build_ts_only_schema(time_index_column_name); // safety: should not failed (UT covers this) - col(time_index_column_name) + ident(time_index_column_name) .cast_to(&DataType::Int64, &input_schema) .unwrap() .cast_to(&DataType::Float64, &input_schema) diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 94769df28a4..709117df294 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -41,7 +41,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, StatisticsArgs, }; use datafusion::prelude::{Column, Expr}; -use datafusion_expr::{EmptyRelation, col}; +use datafusion_expr::{EmptyRelation, ident}; use datatypes::arrow_array::string_array_value_at_index; use datatypes::prelude::{ConcreteDataType, DataType as GtDataType}; use datatypes::value::{OrderedF64, Value, ValueRef}; @@ -141,14 +141,14 @@ impl UserDefinedLogicalNodeCore for HistogramFold { } let mut exprs = vec![ - col(&self.le_column), - col(&self.ts_column), - col(&self.field_column), + ident(&self.le_column), + ident(&self.ts_column), + ident(&self.field_column), ]; exprs.extend(self.input.schema().fields().iter().filter_map(|f| { let name = f.name(); if name != &self.le_column && name != &self.ts_column && name != &self.field_column { - Some(col(name)) + Some(ident(name)) } else { None } @@ -846,6 +846,8 @@ impl HistogramFoldStream { self.find_first_complete_bucket(&batch) } + /// Returns the first group's candidate bucket count once a group boundary is + /// observed. Optimistic validation checks its bounds before folding. fn find_first_complete_bucket(&self, batch: &RecordBatch) -> DataFusionResult> { if batch.num_rows() == 0 { return Ok(None); @@ -853,21 +855,14 @@ impl HistogramFoldStream { let vectors = Helper::try_into_vectors(batch.columns()) .map_err(|e| DataFusionError::Execution(e.to_string()))?; - let le_array = batch.column(self.le_column_index); - let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len()); self.collect_tag_values(&vectors, 0, &mut tag_values_buf); - let mut group_start = 0usize; - for row in 0..batch.num_rows() { + for row in 1..batch.num_rows() { if !self.is_same_group(&vectors, row, &tag_values_buf) { - // new group begins - self.collect_tag_values(&vectors, row, &mut tag_values_buf); - group_start = row; - } - - if Self::is_positive_infinity(le_array, row) { - return Ok(Some(row - group_start + 1)); + // A new series/timestamp proves that the previous group is + // complete, including any numerically equal +Inf boundaries. + return Ok(Some(row)); } } @@ -888,7 +883,10 @@ impl HistogramFoldStream { let field_array = field_array.as_primitive::(); let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len()); - while remaining_rows >= bucket_num && self.mode == FoldMode::Optimistic { + // Keep one row of lookahead: +Inf alone does not prove completion. + // A group ending at the batch boundary stays buffered until more input + // arrives, or flush_remaining completes it at EOF. + while remaining_rows > bucket_num && self.mode == FoldMode::Optimistic { self.collect_tag_values(&vectors, cursor, &mut tag_values_buf); if !self.validate_optimistic_group( &vectors, @@ -987,6 +985,9 @@ impl HistogramFoldStream { tag_values: &[ValueRef<'_>], ) -> bool { let inf_index = cursor + bucket_num - 1; + if self.is_same_group(vectors, inf_index + 1, tag_values) { + return false; + } if !Self::is_positive_infinity(le_array, inf_index) { return false; } @@ -1217,7 +1218,7 @@ impl HistogramFoldStream { fn flush_remaining(&mut self) -> DataFusionResult<()> { if self.mode == FoldMode::Optimistic && self.input_buffered_rows > 0 { - let buffered_batches: Vec<_> = self.input_buffer.drain(..).collect(); + let buffered_batches: Vec<_> = std::mem::take(&mut self.input_buffer); if !buffered_batches.is_empty() { let batch = concat_batches(&self.input_schema, buffered_batches.as_slice())?; self.switch_to_safe_mode(batch)?; @@ -1274,8 +1275,31 @@ impl HistogramFoldStream { } } - // check input value - if !bucket.windows(2).all(|w| w[0] <= w[1]) { + // Detect equal numeric boundaries during validation so unique buckets + // need neither another scan nor an allocation. + let has_duplicates = !bucket.windows(2).all(|w| w[0] < w[1]); + if has_duplicates && !bucket.windows(2).all(|w| w[0] <= w[1]) { + return Ok(f64::NAN); + } + // Counts must be coalesced before repairing monotonicity. + let coalesced = has_duplicates.then(|| { + let mut bounds = Vec::with_capacity(bucket.len()); + let mut counts = Vec::with_capacity(counter.len()); + for (&bound, &count) in bucket.iter().zip(counter) { + if bounds.last() == Some(&bound) { + *counts.last_mut().unwrap() += count; + } else { + bounds.push(bound); + counts.push(count); + } + } + (bounds, counts) + }); + let (bucket, counter) = match &coalesced { + Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), + None => (bucket, counter), + }; + if matches!(operation, HistogramFoldOperation::Quantile(_)) && bucket.len() < 2 { return Ok(f64::NAN); } let counter = match operation { @@ -1339,26 +1363,6 @@ impl HistogramFoldStream { } fn evaluate_fraction(lower: f64, upper: f64, bucket: &[f64], counter: &[f64]) -> f64 { - let coalesced = bucket - .windows(2) - .any(|bounds| bounds[0] == bounds[1]) - .then(|| { - let mut bounds = Vec::with_capacity(bucket.len()); - let mut counts = Vec::with_capacity(counter.len()); - for (&bound, &count) in bucket.iter().zip(counter) { - if bounds.last() == Some(&bound) { - *counts.last_mut().unwrap() += count; - } else { - bounds.push(bound); - counts.push(count); - } - } - (bounds, counts) - }); - let (bucket, counter) = match &coalesced { - Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), - None => (bucket, counter), - }; let total = *counter.last().unwrap(); if total == 0.0 || lower.is_nan() || upper.is_nan() { return f64::NAN; @@ -1466,8 +1470,8 @@ mod test { // 12 items let host_column_1 = Arc::new(StringArray::from(vec![ - "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", - "host_1", "host_1", "host_1", "host_1", + "host_1", "host_1", "host_1", "host_1", "host_1", "host_2", "host_2", "host_2", + "host_2", "host_2", "host_3", "host_3", ])) as _; let le_column_1 = Arc::new(StringArray::from(vec![ "0.001", "0.1", "10", "1000", "+Inf", "0.001", "0.1", "10", "1000", "+inf", "0.001", @@ -1478,14 +1482,14 @@ mod test { ])) as _; // 2 items - let host_column_2 = Arc::new(StringArray::from(vec!["host_1", "host_1"])) as _; + let host_column_2 = Arc::new(StringArray::from(vec!["host_3", "host_3"])) as _; let le_column_2 = Arc::new(StringArray::from(vec!["10", "1000"])) as _; let val_column_2 = Arc::new(Float64Array::from(vec![1.0, 1.0])) as _; // 11 items let host_column_3 = Arc::new(StringArray::from(vec![ - "host_1", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2", - "host_2", "host_2", "host_2", + "host_3", "host_4", "host_4", "host_4", "host_4", "host_4", "host_5", "host_5", + "host_5", "host_5", "host_5", ])) as _; let le_column_3 = Arc::new(StringArray::from(vec![ "+INF", "0.001", "0.1", "10", "1000", "+iNf", "0.001", "0.1", "10", "1000", "+Inf", @@ -1651,10 +1655,10 @@ mod test { | host | val | +--------+-------------------+ | host_1 | 257.5 | -| host_1 | 5.05 | -| host_1 | 0.0004 | -| host_2 | NaN | -| host_2 | 6.040000000000001 | +| host_2 | 5.05 | +| host_3 | 0.0004 | +| host_4 | NaN | +| host_5 | 6.040000000000001 | +--------+-------------------+", ); assert_eq!(result_literal, expected); @@ -1831,7 +1835,7 @@ mod test { Field::new("le", DataType::Utf8, true), Field::new("val", DataType::Float64, true), ])); - let host_column = Arc::new(StringArray::from(vec!["a", "a", "a", "a", "b", "b"])) as _; + let host_column = Arc::new(StringArray::from(vec!["a", "a", "b", "b", "c", "c"])) as _; let le_column = Arc::new(StringArray::from(vec![ "0.1", "+Inf", "0.1", "1.0", "0.1", "+Inf", ])) as _; @@ -1852,8 +1856,8 @@ mod test { | host | val | +------+-----+ | a | 0.1 | -| a | NaN | -| b | 0.1 | +| b | NaN | +| c | 0.1 | +------+-----+", ); assert_eq!(result_literal, expected); @@ -1930,6 +1934,79 @@ mod test { ); } + #[tokio::test] + async fn trailing_invalid_bounds_fall_back_before_eof() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("le", DataType::Utf8, true), + Field::new("val", DataType::Float64, false), + ])); + for trailing_bound in [Some("bad"), None] { + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![ + "a", "a", "a", "b", "b", "b", "c", "c", "c", + ])), + Arc::new(StringArray::from( + [Some("1"), Some("+Inf"), trailing_bound].repeat(3), + )), + Arc::new(Float64Array::from([2.0, 4.0, 99.0].repeat(3))), + ], + ) + .unwrap(); + let fold = build_fold_exec_from_batches(vec![batch.clone()], schema.clone(), 0.5, 0); + let mut stream = HistogramFoldStream { + le_column_index: 1, + field_column_index: 2, + histogram_column_index: None, + operation: fold.operation, + normal_indices: vec![0], + bucket_size: None, + batch_size: 1, + output_schema: fold.output_schema.clone(), + input_schema: schema.clone(), + mode: FoldMode::Optimistic, + safe_group: None, + input_buffer: vec![], + input_buffered_rows: 0, + output_buffer: HistogramFoldStream::empty_output_buffer(&schema).unwrap(), + output_buffered_rows: 0, + input: fold + .input + .execute(0, SessionContext::default().task_ctx()) + .unwrap(), + metric: BaselineMetrics::new(&fold.metric, 0), + }; + + // Neither +Inf nor a batch boundary alone completes the group. + assert!(stream.fold_input(batch.slice(0, 2)).unwrap().is_none()); + assert!(stream.fold_input(batch.slice(2, 1)).unwrap().is_none()); + assert_eq!(stream.mode, FoldMode::Optimistic); + + for (offset, host) in [(3, "a"), (6, "b")] { + let output = stream + .fold_input(batch.slice(offset, 3)) + .unwrap() + .expect("a completed group should produce output before EOF") + .unwrap(); + assert_eq!(stream.mode, FoldMode::Safe); + assert_eq!(stream.input_buffered_rows, 0); + assert!(stream.input_buffer.is_empty()); + assert_eq!(output.num_rows(), 1); + assert_eq!(string_array_value_at_index(output.column(0), 0), Some(host)); + assert_eq!(output.column(1).as_primitive::().value(0), 1.0); + } + + stream.flush_remaining().unwrap(); + let output = stream.take_output_buf().unwrap().unwrap(); + assert_eq!(output.num_rows(), 1); + assert_eq!(string_array_value_at_index(output.column(0), 0), Some("c")); + assert_eq!(output.column(1).as_primitive::().value(0), 1.0); + assert!(stream.take_output_buf().unwrap().is_none()); + } + } + #[tokio::test] async fn safe_mode_handles_misaligned_groups() { let schema = Arc::new(Schema::new(vec![ @@ -2045,6 +2122,171 @@ mod test { assert!((values[1] - 0.55).abs() < 1e-10, "{values:?}"); } + #[test] + fn evaluate_quantile_duplicate_bounds() { + let cases = [ + // Issue #9443: coalesce before repairing the count at 1000. + ( + vec![10.0, 10.0, 100.0, 100.0, 1000.0, f64::INFINITY], + vec![40.0, 40.0, 50.0, 50.0, 50.0, 100.0], + 0.95, + 77.5, + ), + // Repairing the split counts first would incorrectly inflate them. + ( + vec![10.0, 10.0, 100.0, 100.0, f64::INFINITY], + vec![30.0, 10.0, 40.0, 20.0, 60.0], + 0.5, + 7.5, + ), + ( + vec![10.0, f64::INFINITY, f64::INFINITY], + vec![40.0, 20.0, 30.0], + 0.5, + 6.25, + ), + ( + vec![f64::INFINITY, f64::INFINITY], + vec![20.0, 30.0], + 0.5, + f64::NAN, + ), + ( + vec![-0.0, 0.0, 10.0, f64::INFINITY], + vec![10.0, 10.0, 40.0, 40.0], + 0.75, + 5.0, + ), + ]; + for (bounds, counts, quantile, expected) in cases { + let actual = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(quantile.into()), + &bounds, + &counts, + ) + .unwrap(); + if expected.is_nan() { + assert!(actual.is_nan()); + } else { + assert!((actual - expected).abs() < 1e-10, "{actual} != {expected}"); + } + } + } + + #[tokio::test] + async fn duplicate_infinity_bounds_across_batches() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + let duplicate = [("10", 20.0), ("10.0", 20.0), ("+Inf", 25.0), ("+inf", 25.0)]; + let regular = [("10", 20.0), ("10.0", 20.0), ("+Inf", 50.0)]; + // Cover discovery, an already established optimistic bucket count, + // fallback, and the final group at EOF. + for groups in [ + vec![duplicate.as_slice()], + vec![duplicate.as_slice(), regular.as_slice()], + vec![regular.as_slice(), duplicate.as_slice(), regular.as_slice()], + ] { + let mut timestamps = Vec::new(); + let mut bounds = Vec::new(); + let mut counts = Vec::new(); + for (timestamp, group) in groups.iter().enumerate() { + for &(bound, count) in *group { + timestamps.push(timestamp as i64); + bounds.push(bound); + counts.push(count); + } + } + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps)), + Arc::new(StringArray::from(bounds)), + Arc::new(Float64Array::from(counts)), + ], + ) + .unwrap(); + for batch_size in 1..=batch.num_rows() { + let batches = (0..batch.num_rows()) + .step_by(batch_size) + .map(|offset| batch.slice(offset, batch_size.min(batch.num_rows() - offset))) + .collect(); + let fold = build_fold_exec_from_batches(batches, schema.clone(), 0.5, 0); + let result = + datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values: Vec<_> = result + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![6.25; groups.len()], "batch size {batch_size}"); + } + } + } + + #[tokio::test] + async fn duplicate_bounds_across_batches_and_safe_fallback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + // The first group establishes six buckets for optimistic folding. The + // shorter second group forces safe mode; both have duplicate bounds. + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 1000, 1000, 1000, 1000, 1000, 1000, 2000, 2000, 2000, 2000, 2000, + ])), + Arc::new(StringArray::from(vec![ + "10", "10.0", "100", "100.0", "1000", "+Inf", "10", "10.0", "100", "100.0", + "+Inf", + ])), + Arc::new(Float64Array::from(vec![ + 40.0, 40.0, 50.0, 50.0, 50.0, 100.0, 40.0, 40.0, 50.0, 50.0, 100.0, + ])), + ], + ) + .unwrap(); + // Split each pair of equal boundaries across an input batch boundary. + let batches = (0..batch.num_rows()).map(|i| batch.slice(i, 1)).collect(); + let fold = build_fold_exec_from_batches(batches, schema, 0.95, 0); + let result = datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values: Vec<_> = result + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![77.5, 77.5]); + } + #[test] fn evaluate_row_normal_case() { let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY]; diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 45857879d73..4a9a3fe5fb8 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -37,7 +37,7 @@ use datafusion::physical_plan::{ PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, StatisticsArgs, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::compute; use datatypes::timestamp::timestamp_array_to_primitive; use futures::{Stream, StreamExt, ready}; @@ -46,6 +46,7 @@ use prost::Message; use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; +use crate::extension_plan::range_manipulate::{Nanoseconds, ScanNanoseconds, narrow_scan}; use crate::extension_plan::series_divide::SeriesDivide; use crate::extension_plan::{ METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, local_offset, @@ -140,8 +141,8 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { return vec![]; } - let mut exprs = vec![col(&self.time_index_column)]; - exprs.extend(self.staleness_field_columns().map(col)); + let mut exprs = vec![ident(&self.time_index_column)]; + exprs.extend(self.staleness_field_columns().map(ident)); exprs } @@ -691,6 +692,7 @@ impl InstantManipulateStream { let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; + let timestamp_null_count = timestamps.null_count(); let timestamps = timestamps.values(); let len = timestamps.len(); // Shift the native-tick timeline in i128 before comparing samples. Doing @@ -735,28 +737,77 @@ impl InstantManipulateStream { 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 aligned_ts = Vec::with_capacity(estimated_points); - 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; + // Narrow `i64` nanoseconds whenever the batch fits them, exact `i128` otherwise. + let narrow = (timestamp_null_count == 0) + .then(|| { + narrow_scan( + self.time_unit, + self.offset, + self.lookback_delta, + timestamps, + aligned_start, + aligned_end, + ) + }) + .flatten(); + if let Some(scan) = &narrow { + debug_assert_eq!(i64::try_from(first_ns), Ok(scan.shift(timestamps[0]))); + debug_assert_eq!(i64::try_from(last_ns), Ok(scan.shift(timestamps[len - 1]))); + } + let (take_indices, aligned_ts) = match narrow { + Some(scan) => self.scan_instants::( + timestamps, + aligned_start, + aligned_end, + &scan, + estimated_points, + &is_stale, + ), + None => self.scan_instants::( + timestamps, + aligned_start, + aligned_end, + &ScanNanoseconds::::wide(self.time_unit, self.offset, self.lookback_delta), + estimated_points, + &is_stale, + ), + }; + self.take_record_batch_optional(input, take_indices, aligned_ts) + } + + /// Chooses the sample of every aligned evaluation instant, in `Ns` nanoseconds. + /// + /// Keeps the first row among exact timestamp ties and the latest preceding row + /// otherwise; zero lookback admits only exact samples. Staleness is tested after + /// choosing: a selected stale marker suppresses this evaluation rather than + /// falling back to an older finite value. + fn scan_instants( + &self, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + scan: &ScanNanoseconds, + capacity: usize, + is_stale: impl Fn(usize) -> bool, + ) -> (Vec, Vec) { + let len = timestamps.len(); + let mut cursor = 0usize; + let mut take_indices = Vec::with_capacity(capacity); + let mut aligned_ts = Vec::with_capacity(capacity); + for expected_ms in (start..=end).step_by(self.interval as usize) { + let expected = scan.instant(expected_ms); let mut exact_candidate = None; - while cursor < len && to_nanoseconds(timestamps[cursor]) <= expected { - if to_nanoseconds(timestamps[cursor]) == expected && exact_candidate.is_none() { + while cursor < len && scan.shift(timestamps[cursor]) <= expected { + if scan.shift(timestamps[cursor]) == expected && exact_candidate.is_none() { exact_candidate = Some(cursor); } cursor += 1; } - // 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; }; - 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) + let candidate_ts = scan.shift(timestamps[candidate]); + if (candidate_ts == expected || candidate_ts > scan.window_start(expected)) && candidate_ts <= expected && !is_stale(candidate) { @@ -764,7 +815,7 @@ impl InstantManipulateStream { aligned_ts.push(expected_ms); } } - self.take_record_batch_optional(input, take_indices, aligned_ts) + (take_indices, aligned_ts) } /// Helper function to apply "take" on record batch. @@ -1775,6 +1826,81 @@ mod test { .await; } + #[test] + fn nullable_interior_timestamp_matches_wide_scan() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + ), + Field::new("value", DataType::Float64, false), + ])); + let timestamps = [0, i64::MAX, 1]; + let array = TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(NullBuffer::from(vec![false, false, true])), + ); + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(array), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])), + ], + ) + .unwrap(); + let stream = InstantManipulateStream { + offset: 0, + start: 0, + end: 1_000, + lookback_delta: 1_000, + interval: 1_000, + time_index: 0, + time_unit: TimeUnit::Second, + 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, false), + ])), + 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 output_timestamps = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let output_values = output + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let wide = stream.scan_instants::( + ×tamps, + 0, + 1_000, + &ScanNanoseconds::::wide(TimeUnit::Second, 0, 1_000), + 2, + |_| false, + ); + assert_eq!(output_timestamps.values(), &[0]); + assert_eq!(output_values.values(), &[1.0]); + assert_eq!(wide, (vec![0], vec![0])); + assert_eq!(wide.1, output_timestamps.values().to_vec()); + } + #[test] fn exact_ties_select_first_and_lookback_uses_latest() { for (values, expected_timestamp) in [ diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index 1ca015b10f1..b1ea6ca3e7c 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -33,7 +33,7 @@ use datafusion::physical_plan::{ InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, StatisticsArgs, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::array::TimestampMillisecondArray; use datatypes::arrow::datatypes::{SchemaRef, TimestampMillisecondType}; use datatypes::arrow::record_batch::RecordBatch; @@ -93,8 +93,8 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { self.tag_columns .iter() - .map(col) - .chain(std::iter::once(col(&self.time_index_column_name))) + .map(ident) + .chain(std::iter::once(ident(&self.time_index_column_name))) .collect() } diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index 96a0c97778d..38ecb7aa4db 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashSet; +use std::ops::{Add, Mul, Sub}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -20,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::{DataType, Field, SchemaRef, TimeUnit}; +use datafusion::arrow::datatypes::{DataType, Field, Int64Type, SchemaRef, TimeUnit}; use datafusion::arrow::error::ArrowError; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::stats::Precision; @@ -38,7 +39,7 @@ use datafusion::physical_plan::{ InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, StatisticsArgs, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::timestamp::timestamp_array_to_primitive; use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; @@ -318,8 +319,8 @@ impl UserDefinedLogicalNodeCore for RangeManipulate { } let mut exprs = Vec::with_capacity(1 + self.field_columns.len()); - exprs.push(col(&self.time_index)); - exprs.extend(self.field_columns.iter().map(col)); + exprs.push(ident(&self.time_index)); + exprs.extend(self.field_columns.iter().map(ident)); exprs } @@ -702,33 +703,68 @@ impl RangeManipulateStream { new_columns[*index] = new_column; } - // 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(|| { + // The timestamp range payload is always millisecond ABI. A column is reusable as-is + // only when its own data type already is that ABI: the output schema fixes the + // dictionary value type, so a millisecond column carrying a timezone must be + // converted like any other native unit. + let ts_column = input.column(self.time_index); + let timestamp_values: ArrayRef = if self.offset == 0 + && matches!( + ts_column.data_type(), + DataType::Timestamp(TimeUnit::Millisecond, None) + ) { + ts_column.clone() + } else { + // 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(ts_column).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); + // Single pass over the timestamps: an i128 intermediate holds any native + // timestamp shift; only the final millisecond result must fit i64. Null-free + // batches (the dominant PromQL case) take arrow's infallible `unary` kernel, + // which LLVM can vectorize, and report overflow through a shared cell. But + // `unary` visits NULL slots' underlying values, so a NULL slot holding an + // arbitrary overflowing value would raise a false batch-wide error; nullable + // batches therefore take `try_unary`, which applies the closure to valid rows + // only. + let offset_ns = (self.offset as i128) * 1_000_000; + let timestamp_values = if timestamps.null_count() == 0 { + let overflow = std::cell::Cell::new(false); + let shifted = compute::unary::(×tamps, |v| { + let shifted_ns = (v as i128) * scale + offset_ns; + let shifted_ms = shifted_ns / 1_000_000; + match i64::try_from(shifted_ms) { + Ok(ms) => ms, + Err(_) => { + overflow.set(true); + 0 + } + } + }); + if overflow.get() { + return Err(DataFusionError::Execution( + "RangeManipulate timestamp payload overflow".into(), + )); } - 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(|_| { + shifted + } else { + compute::try_unary::(×tamps, |v| { + let shifted_ns = (v as i128) * scale + offset_ns; + i64::try_from(shifted_ns / 1_000_000).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()) + })? + }; + Arc::new(TimestampMillisecondArray::new( + timestamp_values.values().clone(), + timestamp_values.nulls().cloned(), + )) + }; + let ts_range_column = RangeArray::from_ranges(timestamp_values, ranges.clone()) .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))? .into_dict(); new_columns.push(Arc::new(ts_range_column)); @@ -770,6 +806,7 @@ impl RangeManipulateStream { let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; + let timestamp_null_count = timestamps.null_count(); let timestamps = timestamps.values(); let timestamp = |index| (timestamps[index] as i128) * scale + (self.offset as i128) * 1_000_000; @@ -799,36 +836,193 @@ impl RangeManipulateStream { // The intersection is within the declared i64 query bounds. let start = start as i64; let end = end as i64; - let mut ranges = Vec::new(); // 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 as i128) * 1_000_000 - (self.range as i128) * 1_000_000; - - while left < len && timestamp(left) <= start_ts { - left += 1; + // Calculate for every aligned timestamp (`curr_ts`), assuming ordered + // timestamps: the scan compares shifted samples in `i64` nanoseconds + // whenever they fit the batch, and in exact `i128` nanoseconds otherwise. + let narrow = (timestamp_null_count == 0) + .then(|| { + narrow_scan( + self.time_unit, + self.offset, + self.range, + timestamps, + start, + end, + ) + }) + .flatten(); + let ranges = match narrow { + Some(scan) => scan_ranges::(self.interval, timestamps, start, end, &scan), + None => { + let scan = ScanNanoseconds::::wide(self.time_unit, self.offset, self.range); + scan_ranges::(self.interval, timestamps, start, end, &scan) } - right = right.max(left); - while right < len && timestamp(right) <= (curr_ts as i128) * 1_000_000 { - right += 1; - } - - if left == right { - ranges.push((0, 0)); - } else { - ranges.push((left as _, (right - left) as _)); - } - } + }; Ok((ranges, (start, end))) } + + /// Cursor scan of `timestamps` through the narrow `i64` carrier, for benchmarks. + /// + /// `bench_range_fn` times the fast path directly, next to the exact `i128` fallback in + /// [`RangeManipulateStream::bench_scan_wide`]; production selects the carrier through + /// [`narrow_scan`] inside `calculate_range`. Returns `None` when that carrier is + /// rejected, which is exactly the condition that makes the operator fall back to `i128`. + #[doc(hidden)] + pub fn bench_scan_narrow( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Option> { + let scan = narrow_scan(time_unit, offset, window, timestamps, start, end)?; + Some(scan_ranges::(interval, timestamps, start, end, &scan)) + } + + /// Cursor scan of `timestamps` through the exact `i128` carrier, for benchmarks. + /// + /// The carrier is exact for every input, so this times the fallback that + /// [`narrow_scan`] selects whenever the narrow baselines would overflow. + #[doc(hidden)] + pub fn bench_scan_wide( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Vec<(u32, u32)> { + scan_ranges::( + interval, + timestamps, + start, + end, + &ScanNanoseconds::::wide(time_unit, offset, window), + ) + } +} + +/// Cursor scan over the aligned evaluation points, in `Ns` nanoseconds. +fn scan_ranges( + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + scan: &ScanNanoseconds, +) -> Vec<(u32, u32)> { + let len = timestamps.len(); + let mut ranges = Vec::new(); + let mut left = 0usize; + let mut right = 0usize; + for curr_ts in (start..=end).step_by(interval as _) { + let curr_ts_ns = scan.instant(curr_ts); + let start_ts = scan.window_start(curr_ts_ns); + + while left < len && scan.shift(timestamps[left]) <= start_ts { + left += 1; + } + right = right.max(left); + while right < len && scan.shift(timestamps[right]) <= curr_ts_ns { + right += 1; + } + + if left == right { + ranges.push((0, 0)); + } else { + ranges.push((left as _, (right - left) as _)); + } + } + ranges +} + +/// Cursor scan baselines in nanoseconds, shared with `InstantManipulate`. +/// +/// The narrow `i64` baselines come from [`narrow_scan`], the exact ones from +/// [`ScanNanoseconds::wide`]. +#[derive(Debug, Clone, Copy)] +pub(super) struct ScanNanoseconds { + /// Nanoseconds per timestamp tick. + scale: Ns, + /// The offset shift in nanoseconds. + offset_ns: Ns, + /// The range or lookback window in nanoseconds. + window_ns: Ns, +} + +/// Nanoseconds carrier of the cursor scans: `i64` where [`narrow_scan`] proved that +/// the batch fits, `i128` for the exact fallback. +pub(super) trait Nanoseconds: + Copy + PartialOrd + From + Add + Sub + Mul +{ +} + +impl Nanoseconds for T where + T: Copy + PartialOrd + From + Add + Sub + Mul +{ +} + +impl ScanNanoseconds { + /// The `i128` baselines, exact for every input. + pub(super) fn wide(time_unit: TimeUnit, offset: Millisecond, window: Millisecond) -> Self { + Self { + scale: nanoseconds_per_native_tick(time_unit), + offset_ns: i128::from(offset) * 1_000_000, + window_ns: i128::from(window) * 1_000_000, + } + } +} + +impl ScanNanoseconds { + /// Shifted nanoseconds of an input timestamp. + pub(super) fn shift(&self, timestamp: i64) -> Ns { + Ns::from(timestamp) * self.scale + self.offset_ns + } + + /// Nanoseconds of an evaluation instant given in milliseconds. + pub(super) fn instant(&self, millis: Millisecond) -> Ns { + Ns::from(millis) * Ns::from(1_000_000) + } + + /// The exclusive lower boundary of the window ending at `instant`. + pub(super) fn window_start(&self, instant: Ns) -> Ns { + instant - self.window_ns + } +} + +/// The narrow `i64` baselines of a sorted batch, or `None` when they would overflow. +pub(super) fn narrow_scan( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, +) -> Option> { + let scale = i64::try_from(nanoseconds_per_native_tick(time_unit)).ok()?; + let offset_ns = offset.checked_mul(1_000_000)?; + let window_ns = window.checked_mul(1_000_000)?; + // Sorted samples, so the checked endpoints bound every interior shift and product. + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + shift(*timestamps.first()?)?; + shift(*timestamps.last()?)?; + // The checked `[start, end]` bounds every instant and every `instant - window`. + start.checked_mul(1_000_000)?.checked_sub(window_ns)?; + end.checked_mul(1_000_000)?.checked_sub(window_ns)?; + Some(ScanNanoseconds { + scale, + offset_ns, + window_ns, + }) } #[cfg(test)] @@ -859,6 +1053,9 @@ mod test { const TIME_INDEX_COLUMN: &str = "timestamp"; + /// The `(offset, length)` selections of one cursor scan. + type RangeSelections = Vec<(u32, u32)>; + fn project_batch(batch: &RecordBatch, indices: &[usize]) -> RecordBatch { let fields = indices .iter() @@ -1711,6 +1908,36 @@ mod test { stream.calculate_range(&batch).unwrap() } + #[test] + fn calculate_range_nullable_interior_timestamp_matches_wide_scan() { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + )])); + let stream = scan_stream_for_test(TimeUnit::Second, 0, 0, 1_000, 1_000, 1_000); + let timestamps = [0, i64::MAX, 1]; + let array = TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(NullBuffer::from(vec![false, false, true])), + ); + let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap(); + + let (ranges, bounds) = stream.calculate_range(&batch).unwrap(); + let wide = RangeManipulateStream::bench_scan_wide( + TimeUnit::Second, + 0, + 1_000, + 1_000, + ×tamps, + bounds.0, + bounds.1, + ); + assert_eq!(bounds, (0, 1_000)); + assert_eq!(ranges, vec![(0, 1), (0, 0)]); + assert_eq!(ranges, wide); + } + #[test] fn calculate_range_keeps_query_aligned_tail() { let (ranges, bounds) = calculate_range_for_test(4, 94, 30, 15, &[20, 50, 80]); @@ -1954,4 +2181,898 @@ mod test { ); } } + + /// Nanoseconds per native tick, as kept by the narrow carrier. + fn native_scale(time_unit: TimeUnit) -> i64 { + i64::try_from(nanoseconds_per_native_tick(time_unit)).unwrap() + } + + /// Shifted nanoseconds of a sample, computed in the exact `i128` carrier. + fn shifted_nanoseconds(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> i128 { + (timestamp as i128) * nanoseconds_per_native_tick(time_unit) + (offset as i128) * 1_000_000 + } + + /// Shifted milliseconds of the instant grid of a sample, flooring like `calculate_range`. + fn shifted_floor_millis( + time_unit: TimeUnit, + offset: Millisecond, + timestamp: i64, + ) -> Millisecond { + Millisecond::try_from( + shifted_nanoseconds(time_unit, offset, timestamp).div_euclid(1_000_000), + ) + .unwrap() + } + + /// Shifted milliseconds kept by the range payload of a sample. + /// + /// The payload conversion truncates toward zero, so samples below the epoch may keep + /// the millisecond above their own native tick. + fn shifted_millis(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> Millisecond { + Millisecond::try_from(shifted_nanoseconds(time_unit, offset, timestamp) / 1_000_000) + .unwrap() + } + + /// Stream whose cursor scan the differential tests below drive directly. + fn scan_stream_for_test( + time_unit: TimeUnit, + offset: Millisecond, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + ) -> RangeManipulateStream { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + )])); + let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap(); + RangeManipulateStream { + offset, + start, + end, + interval, + range, + time_index: 0, + time_unit, + field_columns: vec![], + aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])), + output_schema: schema, + input: Box::pin(empty_stream), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + } + } + + /// Exact expectations of the cursor scan, formulated outside the production loops. + /// + /// Every evaluation instant keeps the samples of `(instant - range, instant]` in + /// shifted nanoseconds: the lower boundary is exclusive, the upper one inclusive, and + /// native ticks below a millisecond stay distinct. + fn scan_oracle( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + interval: Millisecond, + ) -> Vec<(u32, u32)> { + let window_ns = (range as i128) * 1_000_000; + let shifted = timestamps + .iter() + .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp)) + .collect::>(); + (start..=end) + .step_by(interval as usize) + .map(|instant| { + let instant_ns = (instant as i128) * 1_000_000; + let mut offset = None; + let mut length = 0u32; + for (index, timestamp) in shifted.iter().enumerate() { + if *timestamp > instant_ns - window_ns && *timestamp <= instant_ns { + offset.get_or_insert(index as u32); + length += 1; + } + } + (offset.unwrap_or(0), length) + }) + .collect() + } + + /// Runs the cursor scan of one input through both carriers. + fn scan_both_carriers( + stream: &RangeManipulateStream, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> (Option, RangeSelections) { + let exact = scan_ranges::( + stream.interval, + timestamps, + start, + end, + &ScanNanoseconds::::wide(stream.time_unit, stream.offset, stream.range), + ); + let narrow = narrow_scan( + stream.time_unit, + stream.offset, + stream.range, + timestamps, + start, + end, + ) + .map(|scan| scan_ranges::(stream.interval, timestamps, start, end, &scan)); + // The benchmark hooks must measure the very scans the operator runs. + assert_eq!( + RangeManipulateStream::bench_scan_wide( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + exact, + "the i128 benchmark hook diverged from the cursor scan" + ); + if let Some(narrow) = &narrow { + assert_eq!( + RangeManipulateStream::bench_scan_narrow( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + Some(narrow.clone()), + "the i64 benchmark hook diverged from the cursor scan" + ); + } else { + assert_eq!( + RangeManipulateStream::bench_scan_narrow( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + None, + "the i64 benchmark hook must report the fallback" + ); + } + (narrow, exact) + } + + /// Asserts the narrow carrier was taken and matches the exact scan evaluation by evaluation. + fn assert_narrow_carrier_agrees( + narrow: Option>, + exact: &[(u32, u32)], + context: &str, + ) { + let Some(narrow) = narrow else { + panic!("the i64 carrier was expected, but narrow_scan rejected the input: {context}"); + }; + assert_eq!( + narrow.len(), + exact.len(), + "carriers disagree on the number of evaluations: {context}" + ); + for (evaluation, (narrow_range, exact_range)) in narrow.iter().zip(exact).enumerate() { + assert_eq!( + narrow_range, exact_range, + "carriers disagree on evaluation {evaluation}: {context}" + ); + } + } + + /// Replays the cursor scan in checked `i64` arithmetic. + /// + /// Reports whether the narrow carrier overflows on an operation the scan evaluates, + /// which is what makes the `i128` fallback unavoidable for an input. + fn narrow_carrier_overflows( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + interval: Millisecond, + ) -> bool { + let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else { + return true; + }; + let Some(offset_ns) = offset.checked_mul(1_000_000) else { + return true; + }; + let Some(window_ns) = range.checked_mul(1_000_000) else { + return true; + }; + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + let mut left = 0usize; + let mut right = 0usize; + for instant in (start..=end).step_by(interval as usize) { + let Some(instant_ns) = instant.checked_mul(1_000_000) else { + return true; + }; + let Some(window_start) = instant_ns.checked_sub(window_ns) else { + return true; + }; + while left < timestamps.len() { + let Some(shifted) = shift(timestamps[left]) else { + return true; + }; + if shifted <= window_start { + left += 1; + } else { + break; + } + } + right = right.max(left); + while right < timestamps.len() { + let Some(shifted) = shift(timestamps[right]) else { + return true; + }; + if shifted <= instant_ns { + right += 1; + } else { + break; + } + } + } + false + } + + /// The checked `i64` operation that makes [`narrow_scan`] reject an input, if any. + /// + /// Mirrors the rejection criteria of [`narrow_scan`] so the deterministic cases below + /// can pin which boundary drives the `i128` fallback. + fn narrow_carrier_rejection( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Option<&'static str> { + let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else { + return Some("nanoseconds per native tick do not fit i64"); + }; + let Some(offset_ns) = offset.checked_mul(1_000_000) else { + return Some("offset in nanoseconds does not fit i64"); + }; + let Some(window_ns) = range.checked_mul(1_000_000) else { + return Some("range in nanoseconds does not fit i64"); + }; + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + let (Some(first), Some(last)) = (timestamps.first(), timestamps.last()) else { + return Some("empty batch has no endpoints"); + }; + if shift(*first).is_none() || shift(*last).is_none() { + return Some("sample nanoseconds do not fit i64"); + } + if start + .checked_mul(1_000_000) + .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns)) + .is_none() + { + return Some("start bound in nanoseconds does not fit i64"); + } + if end + .checked_mul(1_000_000) + .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns)) + .is_none() + { + return Some("end bound in nanoseconds do not fit i64"); + } + None + } + + #[test] + fn scan_carriers_agree_on_range_boundaries_within_one_millisecond() { + // Membership is `(lower, upper]` on native ticks: the sample exactly on the lower + // boundary is dropped, the sample exactly on the upper boundary is kept, and the + // three samples sharing one millisecond payload stay distinct and survive together. + let cases = [ + ( + TimeUnit::Microsecond, + vec![1_000_000, 1_000_001, 1_000_999, 1_001_000, 1_001_001], + ), + ( + TimeUnit::Nanosecond, + vec![ + 1_000_000_000, + 1_000_000_001, + 1_000_000_999, + 1_001_000_000, + 1_001_000_001, + ], + ), + ]; + + for (time_unit, timestamps) in cases { + let (start, end, interval, range) = (1_001, 1_001, 1, 1); + let context = format!("unit={time_unit:?}, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range); + + let oracle = scan_oracle(time_unit, 0, range, ×tamps, start, end, interval); + assert_eq!( + oracle, + vec![(1, 3)], + "{context}: (lower, upper] expectations" + ); + + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + assert_narrow_carrier_agrees(narrow, &exact, &context); + + // Two distinct native ticks below the upper boundary share the millisecond + // payload `1_000`; they must not be deduplicated by going through that ABI. + let payloads = timestamps[1..4] + .iter() + .map(|timestamp| shifted_millis(time_unit, 0, *timestamp)) + .collect::>(); + assert_eq!( + payloads, + vec![1_000, 1_000, 1_001], + "{context}: kept samples in one millisecond keep both native ticks" + ); + assert_ne!( + timestamps[1], timestamps[2], + "{context}: the two ticks sharing a payload must stay distinct" + ); + } + } + + #[test] + fn scan_falls_back_to_i128_when_only_the_sample_tail_overflows() { + // `narrow_scan` inspects the sorted endpoints, so a batch whose prefix fits the + // `i64` carrier but whose tail does not has to take the exact carrier. The tail + // sample stays outside every window, so the selections keep the prefix ones. + struct Case { + time_unit: TimeUnit, + narrow_safe: Vec, + overflowing: i64, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + } + + let cases = [ + Case { + time_unit: TimeUnit::Second, + narrow_safe: vec![9_223_372_034, 9_223_372_035, 9_223_372_036], + overflowing: 9_223_372_037, + start: 9_223_372_034_000, + end: 9_223_372_036_000, + interval: 1_000, + range: 1_000, + }, + Case { + time_unit: TimeUnit::Millisecond, + narrow_safe: vec![9_223_372_036_853, 9_223_372_036_854], + overflowing: 9_223_372_036_855, + start: 9_223_372_036_853, + end: 9_223_372_036_854, + interval: 1, + range: 1, + }, + Case { + time_unit: TimeUnit::Microsecond, + narrow_safe: vec![9_223_372_036_853_999, 9_223_372_036_854_000], + overflowing: 9_223_372_036_854_776, + start: 9_223_372_036_854, + end: 9_223_372_036_854, + interval: 1, + range: 1, + }, + ]; + + for case in cases { + let Case { + time_unit, + narrow_safe, + overflowing, + start, + end, + interval, + range, + } = case; + let context = format!("unit={time_unit:?}, timestamps={narrow_safe:?} + {overflowing}"); + let scale = native_scale(time_unit); + // The prefix fits the narrow carrier, the tail sample does not... + assert!( + narrow_safe + .iter() + .all(|timestamp| timestamp.checked_mul(scale).is_some()), + "{context}: prefix samples must fit i64 nanoseconds" + ); + assert!( + overflowing.checked_mul(scale).is_none(), + "{context}: the tail sample must overflow i64 nanoseconds" + ); + // ...even though the millisecond payload of the tail stays legal. + let payload = shifted_millis(time_unit, 0, overflowing); + assert!( + payload.abs() < i64::MAX / 2, + "{context}: the shifted payload {payload} must stay legal" + ); + + let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range); + + let (narrow, prefix_exact) = scan_both_carriers(&stream, &narrow_safe, start, end); + let prefix_oracle = + scan_oracle(time_unit, 0, range, &narrow_safe, start, end, interval); + assert_eq!( + prefix_exact, prefix_oracle, + "{context}: prefix i128 carrier" + ); + assert_narrow_carrier_agrees(narrow, &prefix_exact, &format!("{context} prefix")); + + let mut timestamps = narrow_safe.clone(); + timestamps.push(overflowing); + assert_eq!( + narrow_carrier_rejection(time_unit, 0, range, ×tamps, start, end), + Some("sample nanoseconds do not fit i64"), + "{context}: the overflowing endpoint must reject the narrow carrier" + ); + assert!( + narrow_carrier_overflows(time_unit, 0, range, ×tamps, start, end, interval), + "{context}: the scan evaluates the overflowing tail" + ); + + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert!( + narrow.is_none(), + "{context}: the i128 carrier must be taken" + ); + let oracle = scan_oracle(time_unit, 0, range, ×tamps, start, end, interval); + assert_eq!(exact, oracle, "{context}: exact carrier on the full batch"); + assert_eq!( + exact, prefix_exact, + "{context}: the out-of-window tail keeps the prefix selections" + ); + } + } + + #[test] + fn scan_carriers_split_at_the_nanosecond_boundaries() { + // Nanoseconds are not scaled, so the `i64` carrier is rejected by the query bounds: + // below the millisecond of `i64::MIN` the start bound underflows, while a batch near + // `i64::MAX` still fits. Both keep a legal millisecond payload. + let low = i64::MIN; + let low_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, low); + assert_eq!(low_floor, -9_223_372_036_855); + assert!(low_floor.checked_mul(1_000_000).is_none()); + assert_eq!( + shifted_millis(TimeUnit::Nanosecond, 0, low), + -9_223_372_036_854 + ); + + let timestamps = vec![low, low + 1_000_000]; + let (start, end, interval, range) = (low_floor, low_floor + 1, 1, 1); + let context = format!("nanosecond floor, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range); + assert_eq!( + narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, ×tamps, start, end), + Some("start bound in nanoseconds does not fit i64"), + "{context}" + ); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert!( + narrow.is_none(), + "{context}: the i128 carrier must be taken" + ); + let oracle = scan_oracle( + TimeUnit::Nanosecond, + 0, + range, + ×tamps, + start, + end, + interval, + ); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + + // The ceiling side still fits the narrow carrier, so both carriers have to agree. + let high = i64::MAX - 2_000_000; + let high_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, high); + assert!(high_floor.checked_mul(1_000_000).is_some()); + let timestamps = vec![high, high + 1_000_000, high + 2_000_000]; + let (start, end, interval, range) = (high_floor, high_floor + 2, 1, 1); + let context = format!("nanosecond ceiling, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range); + assert_eq!( + narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, ×tamps, start, end), + None, + "{context}: the ceiling batch still fits the narrow carrier" + ); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + let oracle = scan_oracle( + TimeUnit::Nanosecond, + 0, + range, + ×tamps, + start, + end, + interval, + ); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + assert_narrow_carrier_agrees(narrow, &exact, &context); + } + + #[test] + fn scan_carriers_agree_for_seeded_native_precision_matrix() { + // Seeded matrix over all four native precisions: positive and negative offsets, + // negative times, sub-millisecond samples sharing a millisecond payload, and native + // extremes whose intermediate product overflows while the shifted millisecond stays + // legal. + let units = [ + TimeUnit::Second, + TimeUnit::Millisecond, + TimeUnit::Microsecond, + TimeUnit::Nanosecond, + ]; + // `narrow_scan` rejects the two large offsets: `offset * 1ms` overflows `i64`. + let offsets = [0, 1, -1, 250, -250, 9_223_372_036_855, -9_223_372_036_855]; + let extremes = |time_unit: TimeUnit| match time_unit { + TimeUnit::Second => [-9_223_372_037, 9_223_372_037], + TimeUnit::Millisecond => [-9_223_372_036_855, 9_223_372_036_855], + TimeUnit::Microsecond => [-9_223_372_036_854_776, 9_223_372_036_854_776], + TimeUnit::Nanosecond => [i64::MIN, i64::MAX], + }; + let mut prng = TinyPrng(0x5ca1_ab1e_d1ff_0001); + let mut narrow_cases = 0usize; + let mut wide_cases = 0usize; + let mut skipped = 0usize; + + for case in 0..256usize { + let time_unit = units[case % units.len()]; + let offset = offsets[(case / units.len()) % offsets.len()]; + let [low, high] = extremes(time_unit); + let anchor = match prng.next_i64(0, 3) { + 0 => prng.next_i64(-2_000, 2_000), + 1 => -prng.next_i64(1, 2_000), + 2 => low, + _ => high, + }; + // Samples stay within a few milliseconds, with a native granularity small + // enough that sub-millisecond samples share a millisecond payload. + let step = match time_unit { + TimeUnit::Second => 1, + TimeUnit::Millisecond => 1, + TimeUnit::Microsecond => 1 + prng.next_i64(0, 9), + TimeUnit::Nanosecond => 1 + prng.next_i64(0, 999_999), + }; + let count = 1 + prng.next_i64(0, 4); + let ascending = anchor != high; + let mut timestamps = (0..count) + .map(|index| { + let delta = index * step; + if ascending { + anchor + delta + } else { + anchor - delta + } + }) + .collect::>(); + if !ascending { + timestamps.reverse(); + } + // The review shape keeps a legal millisecond payload; skip the extremes that the + // large offsets push out of the millisecond space. + let nanoseconds = timestamps + .iter() + .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp)) + .collect::>(); + if nanoseconds + .iter() + .any(|value| i64::try_from(value / 1_000_000).is_err()) + { + skipped += 1; + continue; + } + + let first_millis = Millisecond::try_from(nanoseconds[0].div_euclid(1_000_000)).unwrap(); + let interval = prng.next_i64(1, 4); + let points = 1 + prng.next_i64(0, 4); + let start = first_millis - interval; + let end = start + interval * (points - 1); + let range = prng.next_i64(0, 3) * interval; + let context = format!( + "case={case}, unit={time_unit:?}, offset={offset}, start={start}, end={end}, \ + interval={interval}, range={range}, timestamps={timestamps:?}" + ); + + let stream = scan_stream_for_test(time_unit, offset, start, end, interval, range); + let oracle = scan_oracle(time_unit, offset, range, ×tamps, start, end, interval); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + + match ( + &narrow, + narrow_carrier_rejection(time_unit, offset, range, ×tamps, start, end), + ) { + (Some(_), None) => { + narrow_cases += 1; + assert!( + !narrow_carrier_overflows( + time_unit, + offset, + range, + ×tamps, + start, + end, + interval + ), + "{context}: every narrow operation must fit i64" + ); + assert_narrow_carrier_agrees(narrow, &exact, &context); + } + (None, Some(_reason)) => { + wide_cases += 1; + } + (narrow, rejection) => panic!( + "{context}: narrow_scan and its mirror disagree, narrow={}, rejection={rejection:?}", + narrow.is_some() + ), + } + } + + assert!( + narrow_cases > 0 && wide_cases > 0, + "both carriers must be covered: {narrow_cases} narrow, {wide_cases} wide \ + ({skipped} skipped)" + ); + assert!( + skipped < 64, + "the generator skipped too many cases: {skipped}" + ); + } + + /// Builds the single-column timestamp batch the cursor scans consume. + fn timestamp_batch(time_unit: TimeUnit, timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + )])); + let array: ArrayRef = match time_unit { + TimeUnit::Second => Arc::new(TimestampSecondArray::from(timestamps.to_vec())), + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps.to_vec())), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps.to_vec())), + }; + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + #[test] + fn calculate_range_dispatches_carriers_and_agrees_pointwise() { + // Production dispatch, not the hooks: `calculate_range` must scan through the narrow + // carrier while the batch fits and through the exact one once the shifted tail does + // not, and both must select the same ranges for the bounds it derives itself. + let narrow_timestamps = [0i64, 15_000, 30_000]; + let wide_timestamps = [9_223_372_036_853, 9_223_372_036_854, 9_223_372_036_855]; + let cases: [(&str, &[i64], Millisecond, Millisecond, bool); 2] = [ + ("narrow carrier", &narrow_timestamps, 0, 30_000, true), + ( + "i128 fallback", + &wide_timestamps, + wide_timestamps[0], + wide_timestamps[1], + false, + ), + ]; + + for (name, timestamps, query_start, query_end, expect_narrow) in cases { + let (interval, range) = if expect_narrow { + (15_000, 15_000) + } else { + (1, 1) + }; + // `scan_stream_for_test` panics on an empty batch when `narrow_scan` is asked for + // the endpoints, so the batch below is never empty. + let stream = scan_stream_for_test( + TimeUnit::Millisecond, + 0, + query_start, + query_end, + interval, + range, + ); + let batch = timestamp_batch(TimeUnit::Millisecond, timestamps); + let (ranges, (start, end)) = stream.calculate_range(&batch).unwrap(); + + let oracle = scan_oracle( + TimeUnit::Millisecond, + 0, + range, + timestamps, + start, + end, + interval, + ); + assert_eq!(ranges, oracle, "{name}: calculate_range pointwise"); + + let rejection = + narrow_carrier_rejection(TimeUnit::Millisecond, 0, range, timestamps, start, end); + assert_eq!( + rejection.is_none(), + expect_narrow, + "{name}: carrier choice by narrow_scan, rejection={rejection:?}" + ); + let (narrow, exact) = scan_both_carriers(&stream, timestamps, start, end); + assert_eq!(exact, ranges, "{name}: i128 carrier is the oracle"); + assert_eq!(narrow.is_some(), expect_narrow, "{name}: narrow carrier"); + assert_eq!( + RangeManipulateStream::bench_scan_wide( + stream.time_unit, + stream.offset, + stream.range, + interval, + timestamps, + start, + end, + ), + ranges, + "{name}: the benchmark hook matches calculate_range" + ); + } + } + + /// Stream whose single millisecond timestamp column is folded into the range payload. + fn payload_stream( + offset: Millisecond, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + ) -> RangeManipulateStream { + let payload_type = DataType::Timestamp(TimeUnit::Millisecond, None); + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, payload_type.clone(), false), + Field::new( + RangeManipulate::build_timestamp_range_name(TIME_INDEX_COLUMN), + RangeArray::convert_data_type(payload_type), + false, + ), + ])); + let mut stream = + scan_stream_for_test(TimeUnit::Millisecond, offset, start, end, interval, range); + stream.output_schema = schema; + stream + } + + /// Builds the single-column millisecond timestamp batch the payload tests consume. + fn timestamp_batch_with_timezone(timezone: Option<&str>, timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, timezone.map(Arc::from)), + false, + )])); + let array: ArrayRef = Arc::new( + TimestampMillisecondArray::from(timestamps.to_vec()).with_timezone_opt(timezone), + ); + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + /// Millisecond payload values of the range timestamp column of a manipulated batch. + fn payload_millis(batch: &RecordBatch) -> Vec { + let dict = batch + .column(batch.num_columns() - 1) + .as_any() + .downcast_ref::>() + .expect("the timestamp range payload must be a folded dictionary") + .clone(); + let payload = RangeArray::try_new(dict).unwrap(); + assert_eq!( + payload.values().data_type(), + &DataType::Timestamp(TimeUnit::Millisecond, None), + "the payload value type must agree with the output schema" + ); + payload + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[test] + fn manipulate_timezone_time_index_skips_zero_copy_payload() { + // Oracle-verified regression: `Timestamp(Millisecond, Some(tz))` passes a + // `time_unit == Millisecond` check but is not the payload ABI. Cloning such a column + // would stamp a dictionary value type disagreeing with the output schema and + // `RecordBatch::try_new` would fail, so the reuse must be guarded on the data type. + let timestamps = [0i64, 1_000, 2_000, 3_000]; + let stream = payload_stream(0, 0, 3_000, 1_000, 1_000); + + // With offset 0 the data type is the only reason to convert, and the `try_unary` + // path is the identity for a millisecond index. + let zoned = stream + .manipulate(timestamp_batch_with_timezone(Some("+08:00"), ×tamps)) + .expect("a timezone-carrying index must still build the output batch") + .expect("the ranges are not all empty"); + assert_eq!( + payload_millis(&zoned), + timestamps, + "the converted payload must keep the try_unary values" + ); + + // The same logical samples without a timezone do take the zero-copy path; both paths + // must agree on the payload. + let plain = stream + .manipulate(timestamp_batch(TimeUnit::Millisecond, ×tamps)) + .unwrap() + .unwrap(); + assert_eq!( + payload_millis(&plain), + payload_millis(&zoned), + "the zero-copy and the converted path must agree" + ); + } + + /// A NULL slot's underlying value must never trigger the payload overflow + /// error: a nullable batch is converted with arrow's `try_unary`, which applies + /// the closure to valid rows only. The same value as a *valid* sample must still + /// error, through the null-free `unary` + overflow-cell path. + #[test] + fn manipulate_null_slot_underlying_value_does_not_overflow() { + // Second-slot underlying value i64::MAX seconds overflows on conversion, + // but the slot is NULL so the conversion must not see it. offset=1ms + // keeps the batch off the zero-copy path. The query range must be bounded + // and contain the first sample: with `end == 0` no range is selected and + // the conversion (and so the regression) is skipped, while an unbounded + // `end` makes the range scan visit `end - start` evaluation points. + let timestamps = [0i64, i64::MAX]; + let validity = NullBuffer::from(vec![true, false]); + let array: ArrayRef = Arc::new(TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(validity), + )); + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + + let mut stream = payload_stream(1, 0, 1_000, 1, 1_000); + stream.time_unit = TimeUnit::Second; + let output = stream + .manipulate(batch) + .expect("a NULL slot's underlying value must not overflow the payload") + .expect("the first sample (0s + 1ms) is inside the query range"); + assert_eq!( + payload_millis(&output)[0], + 1, + "the valid sample must be converted to its shifted millisecond payload" + ); + + // Same underlying value, but now valid: the conversion must reject it. + let array: ArrayRef = Arc::new(TimestampSecondArray::from(timestamps.to_vec())); + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + let mut stream = payload_stream(1, 0, 1_000, 1, 1_000); + stream.time_unit = TimeUnit::Second; + let error = stream + .manipulate(batch) + .expect_err("a valid overflowing timestamp must still error"); + assert!( + error.to_string().contains("timestamp payload overflow"), + "unexpected error: {error}" + ); + } } diff --git a/src/promql/src/extension_plan/scalar_calculate.rs b/src/promql/src/extension_plan/scalar_calculate.rs index e71333837f3..4ad188a1d4f 100644 --- a/src/promql/src/extension_plan/scalar_calculate.rs +++ b/src/promql/src/extension_plan/scalar_calculate.rs @@ -33,7 +33,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, StatisticsArgs, }; use datafusion::prelude::Expr; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::array::{Array, ArrayRef, Float64Array, TimestampMillisecondArray}; use datatypes::arrow::compute::{CastOptions, cast_with_options, concat_batches}; use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; @@ -282,9 +282,9 @@ impl UserDefinedLogicalNodeCore for ScalarCalculate { self.tag_columns .iter() - .map(col) - .chain(std::iter::once(col(&self.time_index))) - .chain(std::iter::once(col(&self.field_column))) + .map(ident) + .chain(std::iter::once(ident(&self.time_index))) + .chain(std::iter::once(ident(&self.field_column))) .collect() } diff --git a/src/promql/src/extension_plan/series_divide.rs b/src/promql/src/extension_plan/series_divide.rs index bf62f2c3332..bdd90696d26 100644 --- a/src/promql/src/extension_plan/series_divide.rs +++ b/src/promql/src/extension_plan/series_divide.rs @@ -38,7 +38,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::compute; use datatypes::arrow_array::string_array_value_at_index; use datatypes::compute::SortOptions; @@ -186,8 +186,8 @@ impl UserDefinedLogicalNodeCore for SeriesDivide { self.tag_columns .iter() - .map(col) - .chain(std::iter::once(col(&self.time_index_column))) + .map(ident) + .chain(std::iter::once(ident(&self.time_index_column))) .collect() } diff --git a/src/promql/src/extension_plan/union_distinct_on.rs b/src/promql/src/extension_plan/union_distinct_on.rs index 774904a27aa..90d9dda637e 100644 --- a/src/promql/src/extension_plan/union_distinct_on.rs +++ b/src/promql/src/extension_plan/union_distinct_on.rs @@ -34,7 +34,7 @@ use datafusion::physical_plan::{ Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, hash_utils, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::compute; use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; @@ -278,12 +278,12 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { let mut exprs = self .compare_key_indices .iter() - .filter_map(|index| fields.get(*index).map(|field| col(field.name()))) + .filter_map(|index| fields.get(*index).map(|field| ident(field.name()))) .collect::>(); if !self.compare_key_indices.contains(&self.ts_col_idx) && let Some(field) = fields.get(self.ts_col_idx) { - exprs.push(col(field.name())); + exprs.push(ident(field.name())); } exprs } diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index 22cd64b9147..34618f0a953 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -27,6 +27,7 @@ mod resets; mod round; #[cfg(test)] mod test_util; +mod vector_matching; pub use aggr_over_time::{ AbsentOverTime, AvgOverTime, CountOverTime, LastOverTime, MaxOverTime, MinOverTime, @@ -61,6 +62,7 @@ pub use quantile::QuantileOverTime; pub use quantile_aggr::{QUANTILE_NAME, quantile_udaf}; pub use resets::Resets; pub use round::Round; +pub use vector_matching::{MatchGroupViolation, UniqueMatchGroup}; use crate::range_array::RangeArray; diff --git a/src/promql/src/functions/idelta.rs b/src/promql/src/functions/idelta.rs index 35ce44ba270..5734733234d 100644 --- a/src/promql/src/functions/idelta.rs +++ b/src/promql/src/functions/idelta.rs @@ -140,8 +140,13 @@ impl IDelta { 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 sampled_interval_ms = ts_values[last_offset] - ts_values[prev_offset]; + // Sub-millisecond samples round to the same ms; skip instead of dividing by zero. + if sampled_interval_ms == 0 { + result_builder.append_null(); + continue; + } + let sampled_interval = sampled_interval_ms as f64 / 1000.0; let last_value = value_values[value_offset + last_position]; let prev_value = value_values[value_offset + prev_position]; @@ -263,4 +268,30 @@ mod test { vec![Some(3.0), None, None], ); } + + #[test] + fn idelta_skips_samples_in_same_millisecond() { + // The last two samples of the first range share 999ms; the second range is 1ms apart. + let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([ + 999i64, 999, 999, 1000, + ])); + let values_array = Arc::new(Float64Array::from_iter([10.0, 13.0, 999.0, 1000.0])); + let ranges = [(0, 3), (2, 2)]; + + 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![None, Some(1.0)], + ); + + simple_range_udf_runner( + IDelta::::scalar_udf(), + RangeArray::from_ranges(ts_array, ranges).unwrap(), + RangeArray::from_ranges(values_array, ranges).unwrap(), + vec![], + vec![None, Some(1000.0)], + ); + } } diff --git a/src/promql/src/functions/predict_linear.rs b/src/promql/src/functions/predict_linear.rs index dc49ec5d9f9..eb1ac7b5f07 100644 --- a/src/promql/src/functions/predict_linear.rs +++ b/src/promql/src/functions/predict_linear.rs @@ -46,6 +46,8 @@ impl PredictLinear { RangeArray::convert_data_type(DataType::Float64), // t DataType::Int64, + // evaluation timestamp + DataType::Timestamp(TimeUnit::Millisecond, None), ]; create_udf( Self::name(), @@ -58,11 +60,12 @@ impl PredictLinear { fn predict_linear(input: &[ColumnarValue]) -> Result { error::ensure( - input.len() == 3, - DataFusionError::Plan("prom_predict_linear function should have 3 inputs".to_string()), + input.len() == 4, + DataFusionError::Plan("prom_predict_linear function should have 4 inputs".to_string()), )?; let t_col = &input[2]; + let eval_ts_col = &input[3]; let ts_range = extract_range_array(&input[0])?; let value_range = extract_range_array(&input[1])?; @@ -128,6 +131,51 @@ impl PredictLinear { Box::new(t_array.iter()) } }; + // The evaluation instant the window above was folded for. It is the instant the regression + // is centered on, which is *not* the last sample's timestamp: a window may end before the + // step being evaluated, and an `@`-anchored window is replayed at every step while keeping + // its own end. See `PromPlanner::create_range_eval_ts_expr`. + let eval_ts_iter: Box>> = match eval_ts_col { + ColumnarValue::Scalar(eval_ts_scalar) => { + let eval_ts = match eval_ts_scalar { + ScalarValue::TimestampMillisecond(Some(eval_ts), _) => *eval_ts, + // For a NULL or otherwise unusable evaluation timestamp, returns NULL array, + // which conforms to PromQL's behavior. + _ => { + let null_array = Float64Array::new_null(ts_range.len()); + return Ok(ColumnarValue::Array(Arc::new(null_array))); + } + }; + Box::new((0..ts_range.len()).map(move |_| Some(eval_ts))) + } + ColumnarValue::Array(eval_ts_array) => { + error::ensure( + eval_ts_array.len() == ts_range.len(), + DataFusionError::Execution(format!( + "{}: evaluation timestamp array should have the same length as other columns, found {} and {}", + Self::name(), + eval_ts_array.len(), + ts_range.len() + )), + )?; + match eval_ts_array.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, _) => Box::new( + eval_ts_array + .as_any() + .downcast_ref::() + .expect("checked by data type") + .iter(), + ), + other => { + return Err(DataFusionError::Execution(format!( + "{}: expect TimestampMillisecond as evaluation timestamp array's type, found {}", + Self::name(), + other + ))); + } + } + } + }; let all_timestamps = ts_range .values() .as_any() @@ -140,7 +188,12 @@ impl PredictLinear { .downcast_ref::() .unwrap(); let mut result_builder = Float64Builder::with_capacity(ts_range.len()); - for (index, t) in t_iter.enumerate() { + for (index, (t, eval_ts)) in t_iter.zip(eval_ts_iter).enumerate() { + // A step without an evaluation instant has nothing to extrapolate to. + let Some(eval_ts) = eval_ts else { + result_builder.append_null(); + continue; + }; match predict_linear_impl( &ts_range, &value_range, @@ -148,7 +201,7 @@ impl PredictLinear { all_values, index, t.unwrap(), - Self::name(), + eval_ts, )? { Some(value) => result_builder.append_value(value), None => result_builder.append_null(), @@ -167,7 +220,7 @@ fn predict_linear_impl( all_values: &Float64Array, index: usize, t: i64, - func_name: &str, + eval_ts: i64, ) -> Result, DataFusionError> { let (ts_offset, ts_len) = ts_range.get_offset_length(index).unwrap(); let (value_offset, value_len) = value_range.get_offset_length(index).unwrap(); @@ -175,22 +228,28 @@ fn predict_linear_impl( ts_len == value_len, DataFusionError::Execution(format!( "{}: time and value arrays in a group should have the same length, found {} and {}", - func_name, ts_len, value_len + PredictLinear::name(), + ts_len, + value_len )), )?; if ts_len < 2 { return Ok(None); } - // last timestamp is evaluation timestamp - let evaluate_ts = all_timestamps[ts_offset + ts_len - 1]; + // Like Prometheus, the regression is centered on the evaluation timestamp: the returned + // intercept is the value the window's trend predicts at the step being evaluated, and `t` is + // the horizon (in seconds) that is added on top of it. Centering on the last sample instead + // would drop the distance between the window's end and the step, which is non-zero both + // without `@` (the window can end before the step) and with it (the same anchored window is + // replayed at every step). let (slope, intercept) = linear_regression_slices( all_timestamps, ts_offset, all_values, value_offset, value_len, - evaluate_ts, + eval_ts, ); if slope.is_none() || intercept.is_none() { @@ -231,6 +290,11 @@ mod test { (ts_range_array, value_range_array) } + /// The evaluation timestamp at which the window built by [`build_test_range_arrays`] ends. + fn window_end_eval_ts() -> ScalarValue { + ScalarValue::TimestampMillisecond(Some(3000), None) + } + #[test] fn calculate_predict_linear_none() { let ts_array = Arc::new(TimestampMillisecondArray::from_iter( @@ -244,7 +308,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(0))], + vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()], vec![None, None], ); } @@ -256,7 +320,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(0))], + vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()], // value at t = 0 vec![Some(38.63636363636364)], ); @@ -269,7 +333,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(3000))], + vec![ScalarValue::Int64(Some(3000)), window_end_eval_ts()], // value at t = 3000 vec![Some(31856.818181818187)], ); @@ -282,7 +346,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(4200))], + vec![ScalarValue::Int64(Some(4200)), window_end_eval_ts()], // value at t = 4200 vec![Some(44584.09090909091)], ); @@ -295,7 +359,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(6600))], + vec![ScalarValue::Int64(Some(6600)), window_end_eval_ts()], // value at t = 6600 vec![Some(70038.63636363638)], ); @@ -308,7 +372,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(7800))], + vec![ScalarValue::Int64(Some(7800)), window_end_eval_ts()], // value at t = 7800 vec![Some(82765.9090909091)], ); @@ -327,7 +391,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(0))], + vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()], vec![Some(30.0)], ); } @@ -348,9 +412,69 @@ mod test { ColumnarValue::Array(Arc::new(ts_dict)), ColumnarValue::Array(Arc::new(value_dict)), ColumnarValue::Scalar(ScalarValue::Int64(Some(0))), + ColumnarValue::Scalar(window_end_eval_ts()), ]) .unwrap_err(); assert!(err.to_string().contains("Empty range is not expected")); } + + /// The regression is centered on the evaluation timestamp, so the returned intercept follows + /// the step being evaluated even though the window stays the same. + #[test] + fn calculate_predict_linear_centers_on_eval_ts() { + // One sample per second with values 0, 1, 2: the trend is exactly one unit per second. + for (eval_ts, t, expected) in [ + // The window ends at 2s: evaluating there returns its last value. + (2000, 0, 2.0), + // One second past the window's end the trend is at 3. + (3000, 0, 3.0), + // One second past the window's end, predicting another two seconds ahead. + (3000, 2, 5.0), + ] { + let ts_values = Arc::new(TimestampMillisecondArray::from_iter( + [0i64, 1000, 2000].into_iter().map(Some), + )); + let value_values = Arc::new(Float64Array::from_iter([0.0, 1.0, 2.0])); + let ts_array = RangeArray::from_ranges(ts_values, [(0, 3)]).unwrap(); + let value_array = RangeArray::from_ranges(value_values, [(0, 3)]).unwrap(); + + simple_range_udf_runner( + PredictLinear::scalar_udf(), + ts_array, + value_array, + vec![ + ScalarValue::Int64(Some(t)), + ScalarValue::TimestampMillisecond(Some(eval_ts), None), + ], + vec![Some(expected)], + ); + } + } + + /// The evaluation timestamp arrives as a column, so the UDF accepts it as an array too. + #[test] + fn calculate_predict_linear_accepts_eval_ts_array() { + let (ts_array, value_array) = build_test_range_arrays(); + let eval_ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([3000])); + + let result = PredictLinear::predict_linear(&[ + ColumnarValue::Array(Arc::new(ts_array.into_dict())), + ColumnarValue::Array(Arc::new(value_array.into_dict())), + ColumnarValue::Scalar(ScalarValue::Int64(Some(0))), + ColumnarValue::Array(eval_ts_array), + ]) + .unwrap(); + + let result = result + .into_array(1) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(result.len(), 1); + // The same value `calculate_predict_linear_test1` gets from a scalar evaluation timestamp. + assert!((result.value(0) - 38.63636363636364).abs() < 0.0001); + } } diff --git a/src/promql/src/functions/vector_matching.rs b/src/promql/src/functions/vector_matching.rs new file mode 100644 index 00000000000..ef96e8aa5ea --- /dev/null +++ b/src/promql/src/functions/vector_matching.rs @@ -0,0 +1,213 @@ +// 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. + +//! Cardinality checks for PromQL vector matching. + +use datafusion::arrow::array::{Array, Int64Array}; +use datafusion::arrow::util::display::array_value_to_string; +use datafusion::common::{DataFusionError, Result as DfResult}; +use datafusion::logical_expr::{ScalarUDF, Volatility}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_common::ScalarValue; +use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature}; +use datatypes::arrow::datatypes::DataType; + +use crate::functions::extract_array; + +/// The rule a repeated match group violates. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MatchGroupViolation { + /// Several series share a match group on the side that must hold one series per group. + DuplicateOnOneSide { one_side_is_left: bool }, + /// One-to-one matching found several matches for the same group. + ImplicitManyToOne, + /// A group modifier left several matches with the same result label set. + AmbiguousGroupLabels, + /// A function rewrote labels so that several series share the same label set. + DuplicateLabelSet, +} + +impl MatchGroupViolation { + fn message(&self, group: &str) -> String { + match self { + Self::DuplicateOnOneSide { one_side_is_left } => { + let side = if *one_side_is_left { "left" } else { "right" }; + format!( + "found duplicate series for the match group {group} on the {side} hand-side \ + of the operation; many-to-many matching not allowed: matching labels must \ + be unique on one side" + ) + } + Self::ImplicitManyToOne => format!( + "multiple matches for labels {group}: many-to-one matching must be explicit \ + (group_left/group_right)" + ), + Self::AmbiguousGroupLabels => format!( + "multiple matches for labels {group}: grouping labels must ensure unique matches" + ), + Self::DuplicateLabelSet => { + "vector cannot contain metrics with the same labelset".to_string() + } + } + } +} + +/// Rejects a vector matching whose match groups are not unique. +/// +/// Takes the per-group row count as first argument and the label columns that form the group +/// as the remaining ones. Returns `true` when every group holds a single row, and fails the +/// query otherwise; PromQL has no way to express the duplicate series it would produce. +pub struct UniqueMatchGroup; + +impl UniqueMatchGroup { + pub const fn name() -> &'static str { + "prom_assert_unique_match_group" + } + + pub fn scalar_udf(labels: Vec, violation: MatchGroupViolation) -> ScalarUDF { + ScalarUDF::new_from_impl(AssertUniqueMatchGroup { + signature: Signature::variadic_any(Volatility::Volatile), + labels, + violation, + }) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct AssertUniqueMatchGroup { + signature: Signature, + /// Label names of the group, in the same order as the label arguments. + labels: Vec, + violation: MatchGroupViolation, +} + +impl AssertUniqueMatchGroup { + fn render_group(&self, args: &[ColumnarValue], row: usize) -> DfResult { + let mut rendered = Vec::with_capacity(self.labels.len()); + for (label, arg) in self.labels.iter().zip(args) { + let array = extract_array(arg)?; + // A scalar argument was expanded to a single row above. + let row = if array.len() == 1 { 0 } else { row }; + if row >= array.len() || array.is_null(row) { + continue; + } + let value = array_value_to_string(&array, row) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + if value.is_empty() { + continue; + } + rendered.push(format!("{label}=\"{value}\"")); + } + Ok(format!("{{{}}}", rendered.join(", "))) + } +} + +impl ScalarUDFImpl for AssertUniqueMatchGroup { + fn name(&self) -> &str { + UniqueMatchGroup::name() + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { + let Some((counts, labels)) = args.args.split_first() else { + return Err(DataFusionError::Execution(format!( + "{} expects the match group row count as first argument", + UniqueMatchGroup::name() + ))); + }; + let counts = extract_array(counts)?; + let counts = counts + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{} expects an Int64 match group row count, found {}", + UniqueMatchGroup::name(), + counts.data_type() + )) + })?; + + if let Some(row) = + (0..counts.len()).find(|row| !counts.is_null(*row) && counts.value(*row) > 1) + { + let group = self.render_group(labels, row)?; + return Err(DataFusionError::Execution(self.violation.message(&group))); + } + + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::array::StringArray; + use datafusion::arrow::datatypes::Field; + + use super::*; + + fn invoke(counts: Vec, hosts: Vec>) -> DfResult { + let udf = UniqueMatchGroup::scalar_udf( + vec!["host".to_string()], + MatchGroupViolation::ImplicitManyToOne, + ); + let number_rows = counts.len(); + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(counts))), + ColumnarValue::Array(Arc::new(StringArray::from(hosts))), + ], + arg_fields: vec![ + Arc::new(Field::new("count", DataType::Int64, true)), + Arc::new(Field::new("host", DataType::Utf8, true)), + ], + number_rows, + return_field: Arc::new(Field::new("assert", DataType::Boolean, false)), + config_options: Arc::new(Default::default()), + }) + } + + #[test] + fn unique_groups_pass() { + let result = invoke(vec![1, 1], vec![Some("a"), Some("b")]).unwrap(); + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) + )); + } + + #[test] + fn duplicate_group_reports_its_labels() { + let err = invoke(vec![1, 2], vec![Some("a"), Some("b")]).unwrap_err(); + assert!( + err.to_string() + .contains("multiple matches for labels {host=\"b\"}"), + "{err}" + ); + } + + #[test] + fn null_label_is_omitted_from_the_group() { + let err = invoke(vec![2], vec![None]).unwrap_err(); + assert!(err.to_string().contains("labels {}"), "{err}"); + } +} diff --git a/src/puffin/Cargo.toml b/src/puffin/Cargo.toml index 078e3aebeb6..a51ca99d444 100644 --- a/src/puffin/Cargo.toml +++ b/src/puffin/Cargo.toml @@ -34,6 +34,3 @@ snafu.workspace = true tokio.workspace = true tokio-util.workspace = true uuid.workspace = true - -[dev-dependencies] -common-base = { workspace = true, features = ["testing"] } diff --git a/src/puffin/src/lib.rs b/src/puffin/src/lib.rs index 96a8421f98b..2be956e43dc 100644 --- a/src/puffin/src/lib.rs +++ b/src/puffin/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(trait_alias)] - pub mod blob_metadata; pub mod error; pub mod file_format; diff --git a/src/puffin/src/puffin_manager/stager.rs b/src/puffin/src/puffin_manager/stager.rs index 512e94f4e85..8a0adb34bca 100644 --- a/src/puffin/src/puffin_manager/stager.rs +++ b/src/puffin/src/puffin_manager/stager.rs @@ -43,13 +43,13 @@ pub type DirWriterProviderRef = Box; /// /// `Stager` will provide a `BoxWriter` that the caller of `get_blob` /// can use to write the blob into the staging area. -pub trait InitBlobFn = FnOnce(BoxWriter) -> WriteResult; +pub type InitBlobFn<'a> = dyn FnOnce(BoxWriter) -> WriteResult + Send + Sync + 'a; /// Function that initializes a directory. /// /// `Stager` will provide a `DirWriterProvider` that the caller of `get_dir` /// can use to write files inside the directory into the staging area. -pub trait InitDirFn = FnOnce(DirWriterProviderRef) -> WriteResult; +pub type InitDirFn<'a> = dyn FnOnce(DirWriterProviderRef) -> WriteResult + Send + Sync + 'a; /// `Stager` manages the staging area for the puffin files. #[async_trait] @@ -67,7 +67,7 @@ pub trait Stager: Send + Sync { &self, handle: &Self::FileHandle, key: &str, - init_factory: Box, + init_factory: Box>, ) -> Result; /// Retrieves a directory, initializing it if necessary using the provided `init_fn`. @@ -79,7 +79,7 @@ pub trait Stager: Send + Sync { &self, handle: &Self::FileHandle, key: &str, - init_fn: Box, + init_fn: Box>, ) -> Result<(Self::Dir, DirMetrics)>; /// Stores a directory in the staging area. diff --git a/src/puffin/src/puffin_manager/stager/bounded_stager.rs b/src/puffin/src/puffin_manager/stager/bounded_stager.rs index dfb9285452b..e0952b2892e 100644 --- a/src/puffin/src/puffin_manager/stager/bounded_stager.rs +++ b/src/puffin/src/puffin_manager/stager/bounded_stager.rs @@ -146,7 +146,7 @@ impl Stager for BoundedStager { &self, handle: &Self::FileHandle, key: &str, - init_fn: Box, + init_fn: Box>, ) -> Result { let handle_str = handle.to_string(); let cache_key = Self::encode_cache_key(&handle_str, key); @@ -202,7 +202,7 @@ impl Stager for BoundedStager { &self, handle: &Self::FileHandle, key: &str, - init_fn: Box, + init_fn: Box>, ) -> Result<(Self::Dir, DirMetrics)> { let handle_str = handle.to_string(); @@ -332,10 +332,7 @@ impl BoundedStager { BASE64_URL_SAFE.encode(hash) } - async fn write_blob( - target_path: &PathBuf, - init_fn: Box, - ) -> Result { + async fn write_blob(target_path: &PathBuf, init_fn: Box>) -> Result { // To guarantee the atomicity of writing the file, we need to write // the file to a temporary file first... let tmp_path = target_path.with_extension(TMP_EXTENSION); @@ -354,10 +351,7 @@ impl BoundedStager { Ok(size) } - async fn write_dir( - target_path: &PathBuf, - init_fn: Box, - ) -> Result { + async fn write_dir(target_path: &PathBuf, init_fn: Box>) -> Result { // To guarantee the atomicity of writing the directory, we need to write // the directory to a temporary directory first... let tmp_base = target_path.with_extension(TMP_EXTENSION); diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index 480c95a5adb..1d95f7856d5 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -11,12 +11,10 @@ workspace = true default = ["ai_functions"] ai_functions = ["common-function/ai_functions"] enterprise = ["sql/enterprise"] -vector_index = [] [dependencies] ahash.workspace = true api.workspace = true -arc-swap = "1.0" arrow.workspace = true arrow-schema.workspace = true async-recursion = "1.0" @@ -50,11 +48,9 @@ datafusion-physical-expr.workspace = true datafusion-proto.workspace = true datafusion-sql.workspace = true datatypes.workspace = true -either.workspace = true futures.workspace = true futures-util.workspace = true greptime-proto.workspace = true -humantime.workspace = true itertools.workspace = true jsonb.workspace = true lazy_static.workspace = true @@ -63,7 +59,6 @@ meter-core.workspace = true meter-macros.workspace = true object-store.workspace = true once_cell.workspace = true -parking_lot.workspace = true partition.workspace = true prometheus.workspace = true promql.workspace = true @@ -93,13 +88,9 @@ common-query = { workspace = true, features = ["testing"] } fastrand = "2.0" moka = { workspace = true, features = ["future"] } nalgebra.workspace = true -num = "0.4" -num-traits = "0.2" paste.workspace = true pretty_assertions.workspace = true rand.workspace = true -session = { workspace = true, features = ["testing"] } store-api.workspace = true table = { workspace = true, features = ["testing"] } -tokio-stream.workspace = true toml.workspace = true diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index eae46661ad9..ee745673d49 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -31,7 +31,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -45,7 +45,6 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; -use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -80,29 +79,6 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; -// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared -// HTTP/2 connection are pending, trading bounded memory for request liveness. -async fn forward_record_batches( - mut stream: SendableRecordBatchStream, - batch_tx: tokio::sync::mpsc::UnboundedSender>, -) -> Result<()> { - while let Some(batch) = stream.next().await { - match batch.context(CreateRecordBatchSnafu) { - Ok(batch) => { - if batch_tx.send(Ok(batch)).is_err() { - break; - } - tokio::task::yield_now().await; - } - Err(error) => { - let _ = batch_tx.send(Err(error)); - break; - } - } - } - Ok(()) -} - fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -229,7 +205,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let stream = match data { + let mut stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -240,31 +216,32 @@ impl DatafusionQueryEngine { match dml.op { WriteOp::Insert(_) => { - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_insert_record_batches( - batch_rx, - &table_name, - table.schema(), - query_ctx.clone(), - ); - let ((), (rows, cost)) = try_join(producer, consumer).await?; - affected_rows += rows; - insert_cost += cost; + while let Some(batch) = stream.next().await { + let batch = batch.context(CreateRecordBatchSnafu)?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(&table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } } WriteOp::Delete => { - // Keep DELETE on the same producer/consumer schedule as INSERT so the source - // stream can continue draining while mutation RPCs are pending. - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_delete_record_batches( - batch_rx, - &table_name, - &table, - query_ctx.clone(), - ); - let (rows, ()) = try_join(consumer, producer).await?; - affected_rows += rows; + while let Some(batch) = stream.next().await { + let batch = batch.context(CreateRecordBatchSnafu)?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(&table_name, &table, column_vectors, query_ctx.clone()) + .await?; + } } _ => unreachable!("guarded by the 'ensure!' at the beginning"), } @@ -274,53 +251,6 @@ impl DatafusionQueryEngine { )) } - async fn consume_insert_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table_schema: Arc, - query_ctx: QueryContextRef, - ) -> Result<(usize, usize)> { - let mut affected_rows = 0; - let mut insert_cost = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table_schema.clone()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - // We ignore the insert op. - let output = self - .insert(table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - Ok((affected_rows, insert_cost)) - } - - async fn consume_delete_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table: &TableRef, - query_ctx: QueryContextRef, - ) -> Result { - let mut affected_rows = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - affected_rows += self - .delete(table_name, table, column_vectors, query_ctx.clone()) - .await?; - } - Ok(affected_rows) - } - #[tracing::instrument(skip_all)] async fn delete( &self, diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 2c1d066054a..da26d88b433 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -653,7 +653,7 @@ impl PlanRewriter { )) })?; - for (_col_name, alias_set) in part_cols.iter_mut() { + for alias_set in part_cols.values_mut() { let aliased_cols = aliased_columns_for( &alias_set.clone().into_iter().collect(), plan, diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 554628b3de2..694f2bcde59 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -40,6 +40,7 @@ use datafusion_expr::{ }; use datafusion_functions::datetime::date_bin; use datafusion_functions::datetime::expr_fn::now; +use datafusion_functions::unicode::expr_fn::substring; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; use futures::Stream; @@ -528,188 +529,6 @@ fn dictionary_literals_in_values_keep_their_schema() { ); } -#[cfg(feature = "vector_index")] -mod vector_search_tests { - use std::sync::Arc; - - use common_function::function::Function; - use common_function::scalars::udf::create_udf; - use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::{Expr, LogicalPlanBuilder, Signature, Volatility, col, lit}; - use datatypes::schema::{ColumnSchema, SchemaBuilder}; - use store_api::storage::ConcreteDataType; - use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMeta, TableType}; - use table::table::adapter::DfTableProviderAdapter; - use table::{Table, TableRef}; - - use super::*; - use crate::dist_plan::MergeScanLogicalPlan; - - struct TestVectorFunction { - name: &'static str, - signature: Signature, - } - - impl TestVectorFunction { - fn new(name: &'static str) -> Self { - Self { - name, - signature: Signature::any(2, Volatility::Immutable), - } - } - } - - impl std::fmt::Display for TestVectorFunction { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.name) - } - } - - impl Function for TestVectorFunction { - fn name(&self) -> &str { - self.name - } - - fn return_type( - &self, - _input_types: &[datatypes::arrow::datatypes::DataType], - ) -> datafusion_common::Result { - Ok(datatypes::arrow::datatypes::DataType::Float32) - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn invoke_with_args( - &self, - _args: datafusion_expr::ScalarFunctionArgs, - ) -> datafusion_common::Result { - Err(datafusion_common::DataFusionError::Execution( - "test udf should not be invoked".to_string(), - )) - } - } - - fn build_vector_table(table_id: TableId) -> TableRef { - let schema = { - let columns = vec![ - ColumnSchema::new("k0", ConcreteDataType::string_datatype(), true), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new("v", ConcreteDataType::vector_datatype(2), false), - ]; - Arc::new( - SchemaBuilder::try_from_columns(columns) - .unwrap() - .build() - .unwrap(), - ) - }; - - let table_meta = TableMeta { - schema: schema.clone(), - primary_key_indices: vec![0], - value_indices: vec![2], - engine: "test_engine".to_string(), - next_column_id: 3, - options: Default::default(), - created_on: Default::default(), - updated_on: Default::default(), - partition_key_indices: vec![0], - column_ids: vec![0, 1, 2], - }; - - let table_info = TableInfoBuilder::default() - .table_id(table_id) - .name("t".to_string()) - .catalog_name(DEFAULT_CATALOG_NAME) - .schema_name(DEFAULT_SCHEMA_NAME) - .table_version(0) - .table_type(TableType::Base) - .meta(table_meta) - .build() - .unwrap(); - - let data_source = Arc::new(TestDataSource::new(schema)); - Arc::new(Table::new( - Arc::new(table_info), - FilterPushDownType::Unsupported, - data_source, - )) - } - - fn vector_distance_expr() -> Expr { - let udf = create_udf(Arc::new(TestVectorFunction::new("vec_l2sq_distance"))); - Expr::ScalarFunction(ScalarFunction::new_udf( - Arc::new(udf), - vec![ - col("v"), - lit(ScalarValue::Utf8(Some("[1.0, 2.0]".to_string()))), - ], - )) - } - - #[test] - fn vector_search_rewrite_keeps_sort_in_child_plan() { - init_default_ut_logging(); - let table = build_vector_table(0); - let table_source = Arc::new(DefaultTableSource::new(Arc::new( - DfTableProviderAdapter::new(table), - ))); - - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![vector_distance_expr().sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); - - let plan_str = result.to_string(); - assert!(plan_str.contains("MergeSort: vec_l2sq_distance")); - assert!(plan_str.contains("Sort: vec_l2sq_distance")); - assert!(plan_str.contains(MergeScanLogicalPlan::name())); - } - - #[test] - fn vector_search_rewrite_with_filter_keeps_sort_in_child_plan() { - init_default_ut_logging(); - let table = build_vector_table(0); - let table_source = Arc::new(DefaultTableSource::new(Arc::new( - DfTableProviderAdapter::new(table), - ))); - - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .filter(col("k0").eq(lit("hello"))) - .unwrap() - .sort(vec![vector_distance_expr().sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); - - let plan_str = result.to_string(); - assert!(plan_str.contains("MergeSort: vec_l2sq_distance")); - assert!(plan_str.contains("Sort: vec_l2sq_distance")); - assert!(plan_str.contains("Filter: t.k0 = Utf8(\"hello\")")); - assert!(plan_str.contains(MergeScanLogicalPlan::name())); - } -} - fn try_encode_decode_substrait(plan: &LogicalPlan, state: SessionState) { let sub_plan_bytes = substrait::DFLogicalSubstraitConvertor .encode(plan, crate::query_engine::DefaultSerializer) @@ -1183,6 +1002,40 @@ fn expand_proj_alias_aliased_part_col_aggr() { assert_eq!(expected, result.to_string()); } +/// `substr(pk1, 1, 1)` maps rows from different partitions to the same group, so the +/// aggregate has to be merged on the frontend even though it references `pk1`. +#[test] +fn expand_expr_over_part_col_aggr() { + init_default_ut_logging(); + let test_table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(test_table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .aggregate( + vec![substring(col("pk1"), lit(1i64), lit(1i64)), col("pk2")], + vec![min(col("number"))], + ) + .unwrap() + .build() + .unwrap(); + + let config = ConfigOptions::default(); + let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + + let expected = [ + "Projection: substr(t.pk1,Int64(1),Int64(1)), t.pk2, min(t.number)", + " Aggregate: groupBy=[[substr(t.pk1,Int64(1),Int64(1)), t.pk2]], aggr=[[__min_merge(__min_state(t.number)) AS min(t.number)]]", + " MergeScan [is_placeholder=false, remote_input=[", + "Aggregate: groupBy=[[substr(t.pk1, Int64(1), Int64(1)), t.pk2]], aggr=[[__min_state(t.number)]]", + " TableScan: t", + "]]", + ] + .join("\n"); + assert_eq!(expected, result.to_string()); +} + /// notice that step aggr then part col aggr seems impossible as the partition columns for part col aggr /// can't pass through the step aggr without making step aggr also a part col aggr /// so here only test part col aggr -> step aggr case diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index bc41516ba3c..d02b746a8bf 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -16,17 +16,9 @@ use std::collections::HashSet; use std::sync::Arc; use common_function::aggrs::aggr_wrapper::{StateMergeHelper, is_all_aggr_exprs_steppable}; -#[cfg(feature = "vector_index")] -use common_function::scalars::vector::distance::{ - VEC_COS_DISTANCE, VEC_DOT_PRODUCT, VEC_L2SQ_DISTANCE, -}; use common_telemetry::debug; use datafusion::error::Result as DfResult; -#[cfg(feature = "vector_index")] -use datafusion_common::DataFusionError; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -#[cfg(feature = "vector_index")] -use datafusion_expr::Sort; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use promql::extension_plan::{ EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, @@ -37,37 +29,6 @@ use crate::dist_plan::MergeScanLogicalPlan; use crate::dist_plan::analyzer::AliasMapping; use crate::dist_plan::merge_sort::{MergeSortLogicalPlan, merge_sort_transformer}; -#[cfg(feature = "vector_index")] -fn is_vector_sort(sort: &Sort) -> bool { - if sort.expr.len() != 1 { - return false; - } - let sort_expr = &sort.expr[0].expr; - let Expr::ScalarFunction(func) = sort_expr else { - return false; - }; - matches!( - func.name().to_lowercase().as_str(), - VEC_L2SQ_DISTANCE | VEC_COS_DISTANCE | VEC_DOT_PRODUCT - ) -} - -#[cfg(feature = "vector_index")] -fn vector_sort_transformer(plan: &LogicalPlan) -> DfResult { - let LogicalPlan::Sort(sort) = plan else { - return Err(DataFusionError::Internal(format!( - "vector_sort_transformer expects Sort, got {plan}" - ))); - }; - Ok(TransformerAction { - extra_parent_plans: vec![ - MergeSortLogicalPlan::new(sort.input.clone(), sort.expr.clone(), sort.fetch) - .into_logical_plan(), - ], - new_child_plan: Some(LogicalPlan::Sort(sort.clone())), - }) -} - pub struct StepTransformAction { extra_parent_plans: Vec, new_child_plan: Option, @@ -156,7 +117,14 @@ impl Categorizer { LogicalPlan::Filter(filter) => Self::check_expr(&filter.predicate), LogicalPlan::Window(_) => Commutativity::Unimplemented, LogicalPlan::Aggregate(aggr) => { - let is_all_steppable = is_all_aggr_exprs_steppable(&aggr.aggr_expr); + // The state/merge split maps each group expression to one output column, + // which doesn't hold for grouping sets. + let has_grouping_set = aggr + .group_expr + .iter() + .any(|expr| matches!(expr, Expr::GroupingSet(_))); + let is_all_steppable = + !has_grouping_set && is_all_aggr_exprs_steppable(&aggr.aggr_expr); let matches_partition = Self::check_partition(&aggr.group_expr, &partition_cols); if !matches_partition && is_all_steppable { debug!("Plan is steppable: {plan}"); @@ -196,12 +164,6 @@ impl Categorizer { // sort plan needs to consider column priority // Change Sort to MergeSort which assumes the input streams are already sorted hence can be more efficient. - #[cfg(feature = "vector_index")] - if is_vector_sort(_sort) { - return Ok(Commutativity::TransformedCommutative { - transformer: Some(Arc::new(vector_sort_transformer)), - }); - } Commutativity::ConditionalCommutative(Some(Arc::new(merge_sort_transformer))) } LogicalPlan::Join(_) => Commutativity::NonCommutative, @@ -348,26 +310,34 @@ impl Categorizer { /// Return true if the given expr and partition cols satisfied the rule. /// In this case the plan can be treated as fully commutative. /// - /// So only if all partition columns show up in `exprs`, return true. + /// So only if every partition column is itself one of `exprs`, return true. /// Otherwise return false. /// + /// An expression that only references a partition column, like `substr(host, 3, 1)` + /// or `k % 2`, doesn't count: it can put rows from different partitions into the same + /// group. fn check_partition(exprs: &[Expr], partition_cols: &AliasMapping) -> bool { - let mut ref_cols = HashSet::new(); - for expr in exprs { - expr.add_column_refs(&mut ref_cols); - } - let ref_cols = ref_cols - .into_iter() - .map(|c| c.name.clone()) + let group_cols = exprs + .iter() + .filter_map(|expr| { + let mut expr = expr; + while let Expr::Alias(alias) = expr { + expr = &alias.expr; + } + match expr { + Expr::Column(column) => Some(column.name.clone()), + _ => None, + } + }) .collect::>(); for all_alias in partition_cols.values() { let all_alias = all_alias .iter() .map(|c| c.name.clone()) .collect::>(); - // check if ref columns intersect with all alias of partition columns + // check if group columns intersect with all alias of partition columns // is empty, if it's empty, not all partition columns show up in `exprs` - if ref_cols.intersection(&all_alias).count() == 0 { + if group_cols.intersection(&all_alias).count() == 0 { return false; } } diff --git a/src/query/src/dummy_catalog.rs b/src/query/src/dummy_catalog.rs index 2131423e3fe..8f0a2c7ac60 100644 --- a/src/query/src/dummy_catalog.rs +++ b/src/query/src/dummy_catalog.rs @@ -31,15 +31,14 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion_common::DataFusionError; use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; use datatypes::arrow::datatypes::SchemaRef; +use datatypes::data_type::DataType; use datatypes::types::json_type::JsonNativeType; use futures::stream::BoxStream; use session::context::{QueryContext, QueryContextRef}; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; use store_api::region_engine::RegionEngineRef; -use store_api::storage::{ - RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector, VectorSearchRequest, -}; +use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; use table::TableRef; use table::metadata::{TableId, TableInfoRef}; use table::table::adapter::{dictionary_encode_string_columns, supports_pk_dictionary_encoding}; @@ -211,7 +210,13 @@ impl TableProvider for DummyTableProvider { .iter() .map(|e| { // Simple filter on primary key columns are precisely evaluated. - if let Some(simple_filter) = SimpleFilterEvaluator::try_new(e) { + if let Some(simple_filter) = + SimpleFilterEvaluator::try_new_with_column_type(e, &|name| { + self.metadata + .column_by_name(name) + .map(|column| column.column_schema.data_type.as_arrow_type()) + }) + { if self .metadata .column_by_name(simple_filter.column_name()) @@ -281,14 +286,6 @@ impl DummyTableProvider { } } - pub fn with_vector_search_hint(&self, hint: VectorSearchRequest) { - self.scan_request.lock().unwrap().vector_search = Some(hint); - } - - pub fn get_vector_search_hint(&self) -> Option { - self.scan_request.lock().unwrap().vector_search.clone() - } - pub fn with_sequence(&self, sequence: u64) { self.scan_request.lock().unwrap().memtable_max_sequence = Some(sequence); } @@ -656,6 +653,37 @@ mod tests { FLOW_RETURN_REGION_SEQ, FLOW_SINK_TABLE_ID, }; + #[test] + fn test_nullable_label_filter_pushdown() { + use datafusion_expr::{Cast, col, lit}; + use datatypes::arrow::datatypes::DataType as ArrowDataType; + + let provider = crate::optimizer::test_util::mock_table_provider(RegionId::new(1, 1)); + let cast_label = Expr::Cast(Cast::new(Box::new(col("k0")), ArrowDataType::Utf8)); + let cast_field = Expr::Cast(Cast::new(Box::new(col("v0")), ArrowDataType::Utf8)); + let filters = [ + col("k0").is_null().or(col("k0").not_eq(lit("lo"))), + col("k0").is_null().or(cast_label.not_eq(lit("lo"))), + col("k0").is_null().or(col("v0").not_eq(lit(1.0))), + cast_field.eq(lit("1")), + col("k0").in_list(vec![lit(""), lit("lo")], false), + col("k0").in_list(vec![lit(""), lit("lo")], true), + ]; + assert_eq!( + provider + .supports_filters_pushdown(&filters.iter().collect::>()) + .unwrap(), + vec![ + TableProviderFilterPushDown::Exact, + TableProviderFilterPushDown::Exact, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Inexact, + ] + ); + } + fn test_region_id() -> RegionId { RegionId::new(1024, 1) } diff --git a/src/query/src/error.rs b/src/query/src/error.rs index a2946415570..919ef67b94f 100644 --- a/src/query/src/error.rs +++ b/src/query/src/error.rs @@ -331,13 +331,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Failed to get VECTOR index options"))] - GetVectorIndexOptions { - source: datatypes::error::Error, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display( "Column schema mismatch in CTE {}, original: {:?}, expected: {:?}", cte_name, @@ -455,7 +448,6 @@ impl ErrorExt for Error { GetFulltextOptions { source, .. } | GetSkippingIndexOptions { source, .. } - | GetVectorIndexOptions { source, .. } | Datatypes { source, .. } => source.status_code(), } } @@ -481,7 +473,6 @@ impl ErrorExt for Error { TableMutation { source, .. } => source.retry_hint(), GetFulltextOptions { source, .. } | GetSkippingIndexOptions { source, .. } - | GetVectorIndexOptions { source, .. } | Datatypes { source, .. } => source.retry_hint(), _ => RetryHint::NonRetryable, } diff --git a/src/query/src/lib.rs b/src/query/src/lib.rs index 4b2a2d5264d..aa83c8b8f0e 100644 --- a/src/query/src/lib.rs +++ b/src/query/src/lib.rs @@ -12,12 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(int_roundings)] -#![feature(try_blocks)] -#![feature(stmt_expr_attributes)] -#![feature(iterator_try_collect)] -#![feature(box_patterns)] - mod analyze; pub mod datafusion; pub mod dist_plan; diff --git a/src/query/src/log_query/planner.rs b/src/query/src/log_query/planner.rs index 5bc8d034a12..0855974241e 100644 --- a/src/query/src/log_query/planner.rs +++ b/src/query/src/log_query/planner.rs @@ -142,7 +142,7 @@ impl LogQueryPlanner { let exprs = filters .iter() .filter_map(|filter| self.build_filters(filter, schema).transpose()) - .try_collect::>()?; + .collect::, _>>()?; if exprs.is_empty() { Ok(None) } else { @@ -153,7 +153,7 @@ impl LogQueryPlanner { let exprs = filters .iter() .filter_map(|filter| self.build_filters(filter, schema).transpose()) - .try_collect::>()?; + .collect::, _>>()?; if exprs.is_empty() { Ok(None) } else { @@ -191,7 +191,7 @@ impl LogQueryPlanner { self.build_content_filter_with_expr(col_expr.clone(), filter, &df_schema) .transpose() }) - .try_collect::>()?; + .collect::, _>>()?; if filter_exprs.is_empty() { return Ok(Some(col_expr.is_true())); @@ -285,7 +285,7 @@ impl LogQueryPlanner { self.build_content_filter_with_expr(col_expr.clone(), filter, schema) .transpose() }) - .try_collect::>()?; + .collect::, _>>()?; if exprs.is_empty() { return Ok(None); @@ -326,19 +326,19 @@ impl LogQueryPlanner { let args = args .iter() .map(|expr| self.log_expr_to_df_expr(expr, schema)) - .try_collect::>()?; + .collect::, _>>()?; if let Some(alias) = alias { Ok(aggr_fn.call(args).alias(alias)) } else { Ok(aggr_fn.call(args)) } }) - .try_collect::>()?; + .collect::, _>>()?; let group_exprs = by .iter() .map(|expr| self.log_expr_to_df_expr(expr, schema)) - .try_collect::>()?; + .collect::, _>>()?; Ok((aggr_expr, group_exprs)) } @@ -380,7 +380,7 @@ impl LogQueryPlanner { let args = args .iter() .map(|expr| self.log_expr_to_df_expr(expr, schema)) - .try_collect::>()?; + .collect::, _>>()?; let func = self.session_state.scalar_functions().get(name).context( UnknownScalarFunctionSnafu { name: name.to_string(), diff --git a/src/query/src/optimizer/global_limit.rs b/src/query/src/optimizer/global_limit.rs index 5abb9953b7b..b29907f01ab 100644 --- a/src/query/src/optimizer/global_limit.rs +++ b/src/query/src/optimizer/global_limit.rs @@ -254,22 +254,215 @@ fn inherited_partitioning_to_restore( #[cfg(test)] mod tests { - use datafusion::arrow::array::Int32Array; + use datafusion::arrow::array::{Array, Int32Array}; use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; + use datafusion::execution::TaskContext; use datafusion::physical_expr::expressions::{col, lit}; + use datafusion::physical_optimizer::optimizer::PhysicalOptimizer; + use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, + }; use datafusion::physical_plan::filter::FilterExecBuilder; use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; - use datafusion::physical_plan::limit::GlobalLimitExec; + use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::test::TestMemoryExec; use datafusion_common::{JoinType, NullEquality}; use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr}; use super::*; + async fn optimize_and_collect_twice( + mut plan: Arc, + config: &ConfigOptions, + ) -> Vec> { + let mut results = Vec::with_capacity(2); + for _ in 0..2 { + for rule in PhysicalOptimizer::new().rules { + plan = rule.optimize(plan, config).unwrap(); + } + let batches = datafusion::physical_plan::collect( + Arc::clone(&plan), + Arc::new(TaskContext::default()), + ) + .await + .unwrap(); + results.push( + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(), + ); + } + results + } + + #[tokio::test] + async fn physical_optimizer_keeps_global_distinct_limit_across_two_passes() { + for soft_limit in [false, true] { + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + config.optimizer.enable_distinct_aggregation_soft_limit = soft_limit; + + for mode in [ + AggregateMode::FinalPartitioned, + AggregateMode::SinglePartitioned, + ] { + // The query limit is the only initial limit. + let input = input_with_all_hash_partitions(); + let aggregate = match mode { + AggregateMode::FinalPartitioned => { + agg(agg(input, AggregateMode::Partial), mode) + } + AggregateMode::SinglePartitioned => agg(hash_repartition(input), mode), + _ => unreachable!(), + }; + let plan = + Arc::new(GlobalLimitExec::new(aggregate, 0, Some(1))) as Arc; + + let results = optimize_and_collect_twice(plan, &config).await; + assert_eq!( + results.iter().map(Vec::len).collect::>(), + vec![1, 1], + "soft limit enabled: {soft_limit}, mode: {mode:?}", + ); + } + } + } + + #[tokio::test] + async fn physical_optimizer_keeps_count_over_distinct_limit_across_two_passes() { + use datafusion::datasource::MemTable; + use datafusion::execution::context::{SessionConfig, SessionContext as DFSessionContext}; + use datafusion_common::ScalarValue; + + for soft_limit in [false, true] { + let mut config = SessionConfig::new().with_target_partitions(3); + config + .options_mut() + .optimizer + .enable_distinct_aggregation_soft_limit = soft_limit; + let ctx = DFSessionContext::new_with_config(config.clone()); + let schema = schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from((0..100).collect::>()))], + ) + .unwrap(); + let table = MemTable::try_new(schema, vec![vec![batch]; 3]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + let mut plan = ctx + .sql("SELECT COUNT(*) FROM (SELECT DISTINCT a FROM t LIMIT 1) AS limited") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + + for pass in 0..=2 { + if pass > 0 { + for rule in PhysicalOptimizer::new().rules { + plan = rule.optimize(plan, config.options()).unwrap(); + } + } + let batches = datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()).map(|row| { + ScalarValue::try_from_array(batch.column(0).as_ref(), row).unwrap() + }) + }) + .collect::>(); + assert_eq!( + values, + vec![ScalarValue::Int64(Some(1))], + "soft limit enabled: {soft_limit}, additional optimizer passes: {pass}", + ); + } + } + } + + #[tokio::test] + async fn physical_optimizer_keeps_local_limit_per_partition_across_two_passes() { + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + let input = input_with_all_hash_partitions(); + let local_limit = Arc::new(LocalLimitExec::new( + agg(hash_repartition(input), AggregateMode::SinglePartitioned), + 1, + )) as Arc; + + let results = optimize_and_collect_twice(local_limit, &config).await; + assert_eq!(results.iter().map(Vec::len).collect::>(), vec![3, 3]); + } + + #[tokio::test] + async fn physical_optimizer_keeps_ordered_topk_across_two_passes() { + let schema = schema(); + let losing = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![100]))]) + .unwrap(); + let winning = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + let input = Arc::new( + TestMemoryExec::try_new(&[vec![losing], vec![winning]], schema.clone(), None).unwrap(), + ); + let ordering = ordering(schema.as_ref(), false); + let topk = Arc::new( + SortExec::new( + ordering, + agg_with_limit_options( + input, + AggregateMode::FinalPartitioned, + Some(LimitOptions::new_with_order(1, false)), + ), + ) + .with_fetch(Some(1)), + ) as Arc; + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + assert_eq!( + optimize_and_collect_twice(topk, &config).await, + vec![vec![1], vec![1]], + ); + } + + #[tokio::test] + async fn physical_optimizer_keeps_ordered_offset_limit_across_two_passes() { + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + let offset = Arc::new(GlobalLimitExec::new( + Arc::new(SortExec::new( + ordering(schema().as_ref(), false), + unordered_input(), + )), + 3, + Some(1), + )) as Arc; + + assert_eq!( + optimize_and_collect_twice(offset, &config).await, + vec![vec![2], vec![2]], + ); + } + #[test] fn adds_global_limit_for_multi_partition_filter_fetch() { let filter = filter_fetch(unordered_input(), 1); @@ -563,6 +756,38 @@ mod tests { Arc::new(TestMemoryExec::try_new(&partitions, schema, None).unwrap()) } + fn input_with_all_hash_partitions() -> Arc { + let schema = schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from((0..100).collect::>()))], + ) + .unwrap(); + let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]]; + Arc::new(TestMemoryExec::try_new(&partitions, schema, None).unwrap()) + } + + fn agg(input: Arc, mode: AggregateMode) -> Arc { + agg_with_limit_options(input, mode, None) + } + + fn agg_with_limit_options( + input: Arc, + mode: AggregateMode, + limit_options: Option, + ) -> Arc { + let schema = input.schema(); + let group_by = PhysicalGroupBy::new_single(vec![( + col("a", schema.as_ref()).unwrap(), + "a".to_string(), + )]); + Arc::new( + AggregateExec::try_new(mode, group_by, vec![], vec![], input, schema) + .unwrap() + .with_limit_options(limit_options), + ) + } + fn ordered_input() -> (Arc, LexOrdering) { let schema = schema(); let ordering = ordering(schema.as_ref(), false); diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index edb73559139..4685889c998 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/src/query/src/optimizer/scan_hint.rs @@ -31,10 +31,6 @@ use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector}; use crate::dummy_catalog::DummyTableProvider; -#[cfg(feature = "vector_index")] -mod vector_search; -#[cfg(feature = "vector_index")] -use vector_search::VectorSearchState; /// This rule will traverse the plan to collect necessary hints for leaf /// table scan node and set them in [`ScanRequest`]. Hints include: @@ -93,24 +89,18 @@ impl ScanHintRule { }; let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row; - #[cfg(feature = "vector_index")] - let has_vector_hint = rewriter.vector_search.need_rewrite(); - #[cfg(not(feature = "vector_index"))] - let has_vector_hint = false; - let has_hint = rewriter.order_expr.is_some() - || rewriter.ts_row_selector.is_some() - || use_last_row - || has_vector_hint; + let has_hint = + rewriter.order_expr.is_some() || rewriter.ts_row_selector.is_some() || use_last_row; if !has_hint { return Ok(Transformed::no(LogicalPlan::TableScan(table_scan))); } // A provider can be used by several TableScan nodes. Fork its request // for every hinted use-site before applying hints, rather than mutating - // the shared catalog provider. This keeps order/vector/legacy hints + // the shared catalog provider. This keeps order and aggregate hints // local as well as the new LastRow hint. let adapter = original.clone_for_scan(); - Self::apply_hints(&adapter, rewriter, &table_scan, use_last_row); + Self::apply_hints(&adapter, rewriter); if use_last_row { // Apply the instant-derived hint after the aggregate hint. Both // select LastRow today, and this ordering preserves the existing @@ -156,35 +146,13 @@ impl ScanHintRule { true } - fn apply_hints( - adapter: &DummyTableProvider, - rewriter: &mut ScanHintRewriter, - table_scan: &datafusion_expr::logical_plan::TableScan, - use_last_row: bool, - ) { - #[cfg(not(feature = "vector_index"))] - let _ = (table_scan, use_last_row); + fn apply_hints(adapter: &DummyTableProvider, rewriter: &mut ScanHintRewriter) { if let Some(order_expr) = &rewriter.order_expr { Self::set_order_hint(adapter, order_expr); } if let Some((group_by_cols, order_by_col)) = &rewriter.ts_row_selector { Self::set_time_series_row_selector_hint(adapter, group_by_cols, order_by_col); } - #[cfg(feature = "vector_index")] - if use_last_row { - // LastRow and vector search are mutually exclusive for one scan: - // vector search would bypass the ordinary sort/limit path needed by - // the single-evaluation semantics. Still consume the queued hint so - // it cannot be applied to a later scan of the same table. - let _ = rewriter - .vector_search - .take_vector_request_from_dummy(adapter, &table_scan.table_name); - } else if let Some(vector_request) = rewriter - .vector_search - .take_vector_request_from_dummy(adapter, &table_scan.table_name) - { - adapter.with_vector_search_hint(vector_request); - } } fn set_order_hint(adapter: &DummyTableProvider, order_expr: &Vec) { @@ -307,8 +275,6 @@ struct ScanHintRewriter { ts_stack: Vec, Column)>>, inside_single_evaluation: bool, single_evaluation_stack: Vec, - #[cfg(feature = "vector_index")] - vector_search: VectorSearchState, } impl TreeNodeRewriter for ScanHintRewriter { @@ -355,44 +321,10 @@ impl TreeNodeRewriter for ScanHintRewriter { } } - #[cfg(feature = "vector_index")] - { - if let LogicalPlan::Limit(limit) = &node { - self.vector_search.on_limit_enter(limit); - } - if let LogicalPlan::Sort(sort) = &node { - self.vector_search.on_sort_enter(sort); - } - if is_branching_for_vector(&node) { - self.vector_search.on_branching_enter(); - } - if let LogicalPlan::Filter(filter) = &node { - self.vector_search.on_filter_enter(&filter.predicate); - } - if let LogicalPlan::TableScan(table_scan) = &node { - self.vector_search.on_table_scan(table_scan); - } - } - ScanHintRule::set_hints(node, self) } fn f_up(&mut self, node: LogicalPlan) -> Result> { - #[cfg(feature = "vector_index")] - { - match &node { - LogicalPlan::Limit(_) => self.vector_search.on_limit_exit(), - LogicalPlan::Sort(_) => self.vector_search.on_sort_exit(), - LogicalPlan::Filter(_) => self.vector_search.on_filter_exit(), - LogicalPlan::Subquery(_) | LogicalPlan::SubqueryAlias(_) - if is_branching_for_vector(&node) => - { - self.vector_search.on_branching_exit() - } - _ if node.inputs().len() > 1 => self.vector_search.on_branching_exit(), - _ => {} - } - } if let Some(previous) = self.order_stack.pop() { self.order_expr = previous; } @@ -510,43 +442,6 @@ impl ScanHintRewriter { } } -#[cfg(feature = "vector_index")] -fn is_branching_for_vector(node: &LogicalPlan) -> bool { - if node.inputs().len() > 1 { - return true; - } - - match node { - LogicalPlan::Subquery(subquery) => has_non_inlineable_ops(subquery.subquery.as_ref()), - LogicalPlan::SubqueryAlias(alias) => has_non_inlineable_ops(alias.input.as_ref()), - _ => false, - } -} - -#[cfg(feature = "vector_index")] -fn has_non_inlineable_ops(plan: &LogicalPlan) -> bool { - if matches!( - plan, - LogicalPlan::Limit(_) - | LogicalPlan::Sort(_) - | LogicalPlan::Distinct(_) - | LogicalPlan::Aggregate(_) - | LogicalPlan::Window(_) - | LogicalPlan::Union(_) - | LogicalPlan::Join(_) - ) { - return true; - } - - for input in plan.inputs() { - if has_non_inlineable_ops(input) { - return true; - } - } - - false -} - #[cfg(test)] mod test { use std::collections::HashMap; diff --git a/src/query/src/optimizer/scan_hint/vector_search.rs b/src/query/src/optimizer/scan_hint/vector_search.rs deleted file mode 100644 index 90f45378ecd..00000000000 --- a/src/query/src/optimizer/scan_hint/vector_search.rs +++ /dev/null @@ -1,1066 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::{HashMap, HashSet, VecDeque}; - -use common_function::scalars::vector::distance::{ - VEC_COS_DISTANCE, VEC_DOT_PRODUCT, VEC_L2SQ_DISTANCE, -}; -use common_telemetry::debug; -use datafusion_common::{ScalarValue, TableReference}; -use datafusion_expr::logical_plan::FetchType; -use datafusion_expr::utils::split_conjunction; -use datafusion_expr::{Expr, SortExpr}; -use datatypes::types::parse_string_to_vector_type_value; -use store_api::storage::{VectorDistanceMetric, VectorSearchRequest}; - -use crate::dummy_catalog::DummyTableProvider; - -/// Tracks vector search hints while traversing the logical plan. -/// -/// Vector search requests are emitted only when: -/// - ORDER BY uses a supported vector distance function and its direction matches the metric. -/// - A LIMIT (or Sort.fetch) is present to derive k. -/// - The hint stays within a single input chain (not across join/subquery branches). -/// - The target column is non-nullable, or an explicit IS NOT NULL filter exists. -/// -/// Known limitations: -/// - Dynamic overfetching is not implemented yet. When filters exist or ORDER BY includes -/// additional tie-breaker columns (e.g., ORDER BY distance, id), the current fixed k may -/// return incorrect results. A future improvement should dynamically adjust k based on -/// filter selectivity and secondary sort requirements. -/// - Hints only block at subquery boundaries when the subquery contains non-inlineable -/// operators (Limit/Sort/Distinct/Aggregate/Window). Simple subqueries without these -/// operators allow hints to propagate through. In distributed mode, the dist analyzer -/// may inline subqueries before this rule runs, further reducing isolation. -#[derive(Default)] -pub(crate) struct VectorSearchState { - current_distance: Option, - current_limit: Option, - distance_stack: Vec>, - limit_stack: Vec>, - non_null_columns: HashSet, - non_null_stack: Vec>, - vector_hints: HashMap>, -} - -#[derive(Clone)] -struct VectorDistanceInfo { - table_reference: Option, - column_name: String, - query_vector: Vec, - metric: VectorDistanceMetric, -} - -#[derive(Clone)] -struct VectorLimitInfo { - fetch: usize, - skip: usize, -} - -impl VectorLimitInfo { - fn k(&self) -> Option { - self.fetch.checked_add(self.skip) - } -} - -#[derive(Clone)] -struct VectorHintEntry { - distance: VectorDistanceInfo, - limit: VectorLimitInfo, - non_null_constraint: bool, -} - -#[derive(Clone, Hash, PartialEq, Eq)] -struct ColumnKey { - table_reference: Option, - column_name: String, -} - -impl VectorSearchState { - pub(crate) fn need_rewrite(&self) -> bool { - !self.vector_hints.is_empty() - } - - pub(crate) fn on_branching_enter(&mut self) { - // Clear per-branch state so hints are only derived within a single input chain. - self.distance_stack.push(self.current_distance.take()); - self.limit_stack.push(self.current_limit.take()); - self.non_null_stack - .push(std::mem::take(&mut self.non_null_columns)); - } - - pub(crate) fn on_branching_exit(&mut self) { - // Restore the prior chain state after leaving the branch. - if let Some(previous) = self.limit_stack.pop() { - self.current_limit = previous; - } - if let Some(previous) = self.distance_stack.pop() { - self.current_distance = previous; - } - if let Some(previous) = self.non_null_stack.pop() { - self.non_null_columns = previous; - } - } - - pub(crate) fn on_limit_enter(&mut self, limit: &datafusion_expr::logical_plan::Limit) { - self.limit_stack.push(self.current_limit.take()); - self.current_limit = Self::extract_limit_info(limit); - } - - pub(crate) fn on_limit_exit(&mut self) { - if let Some(previous) = self.limit_stack.pop() { - self.current_limit = previous; - } - } - - pub(crate) fn on_sort_enter(&mut self, sort: &datafusion_expr::logical_plan::Sort) { - // Distance is scoped to the nearest sort, while limit may be inherited from parents. - self.distance_stack.push(self.current_distance.take()); - self.limit_stack.push(self.current_limit.clone()); - let distance = Self::extract_distance_from_sort(sort); - self.current_distance = distance.clone(); - // Sort.fetch is a TopK limit, so we can infer k when no LIMIT is present. - if self.current_limit.is_none() { - self.current_limit = distance - .as_ref() - .and_then(|_| Self::extract_limit_from_sort(sort)); - } - } - - pub(crate) fn on_sort_exit(&mut self) { - if let Some(previous) = self.limit_stack.pop() { - self.current_limit = previous; - } - if let Some(previous) = self.distance_stack.pop() { - self.current_distance = previous; - } - } - - pub(crate) fn on_table_scan(&mut self, table_scan: &datafusion_expr::logical_plan::TableScan) { - self.record_vector_hint(table_scan); - } - - pub(crate) fn on_filter_enter(&mut self, predicate: &Expr) { - self.non_null_stack.push(self.non_null_columns.clone()); - for expr in split_conjunction(predicate) { - if let Expr::IsNotNull(inner) = expr - && let Expr::Column(col) = inner.as_ref() - { - self.non_null_columns.insert(ColumnKey { - table_reference: col.relation.clone(), - column_name: col.name.clone(), - }); - } - } - // TODO: detect non-null constraints from more complex predicates (casts/functions). - } - - pub(crate) fn on_filter_exit(&mut self) { - if let Some(previous) = self.non_null_stack.pop() { - self.non_null_columns = previous; - } - } - - pub(crate) fn take_vector_request_from_dummy( - &mut self, - provider: &DummyTableProvider, - table_name: &TableReference, - ) -> Option { - let hint = self.take_vector_hint(table_name)?; - self.build_vector_request_from_dummy(provider, table_name, &hint) - } - - fn build_vector_request_from_dummy( - &self, - provider: &DummyTableProvider, - table_name: &TableReference, - hint: &VectorHintEntry, - ) -> Option { - let info = &hint.distance; - let k = hint.limit.k()?; - - if let Some(ref hint_table) = info.table_reference - && table_name != hint_table - { - return None; - } - - let metadata = provider.region_metadata(); - let column = metadata.column_by_name(&info.column_name)?; - if column.column_schema.is_nullable() && !hint.non_null_constraint { - debug!( - "Skip vector hint: column '{}' is nullable without IS NOT NULL filter", - info.column_name - ); - return None; - } - - Some(VectorSearchRequest { - column_id: column.column_id, - query_vector: info.query_vector.clone(), - k, - metric: info.metric, - }) - } - - fn extract_distance_info(expr: &Expr) -> Option { - let Expr::ScalarFunction(func) = expr else { - return None; - }; - - let func_name = func.name().to_lowercase(); - let metric = match func_name.as_str() { - VEC_L2SQ_DISTANCE => VectorDistanceMetric::L2sq, - VEC_COS_DISTANCE => VectorDistanceMetric::Cosine, - VEC_DOT_PRODUCT => VectorDistanceMetric::InnerProduct, - _ => return None, - }; - - if func.args.len() != 2 { - return None; - } - - let (table_reference, column_name) = match &func.args[0] { - Expr::Column(col) => (col.relation.clone(), col.name.clone()), - _ => return None, - }; - - let query_vector = Self::extract_query_vector(&func.args[1])?; - - Some(VectorDistanceInfo { - table_reference, - column_name, - query_vector, - metric, - }) - } - - fn extract_distance_from_sort( - sort: &datafusion_expr::logical_plan::Sort, - ) -> Option { - if sort.expr.is_empty() { - debug!("Skip vector hint: Sort has no expressions"); - return None; - } - let sort_expr: &SortExpr = &sort.expr[0]; - let info = Self::extract_distance_info(&sort_expr.expr)?; - let expected_asc = info.metric != VectorDistanceMetric::InnerProduct; - if sort_expr.asc != expected_asc { - return None; - } - - if Self::tie_breakers_allowed(&sort.expr[1..], &info) { - Some(info) - } else { - if sort.expr.len() > 1 { - debug!( - "Skip vector hint: Sort has unsupported tie-breakers ({} expressions)", - sort.expr.len() - ); - } - None - } - } - - fn tie_breakers_allowed(sort_exprs: &[SortExpr], distance_info: &VectorDistanceInfo) -> bool { - if sort_exprs.is_empty() { - return true; - } - - sort_exprs.iter().all(|sort_expr| { - let Expr::Column(col) = &sort_expr.expr else { - return false; - }; - - match &distance_info.table_reference { - Some(table) => col.relation.as_ref() == Some(table), - None => col.relation.is_none(), - } - }) - } - - fn extract_limit_info(limit: &datafusion_expr::logical_plan::Limit) -> Option { - let fetch = match limit.get_fetch_type().ok()? { - FetchType::Literal(fetch) => fetch?, - FetchType::UnsupportedExpr => return None, - }; - let skip = match limit.get_skip_type().ok()? { - datafusion_expr::logical_plan::SkipType::Literal(skip) => skip, - datafusion_expr::logical_plan::SkipType::UnsupportedExpr => return None, - }; - Some(VectorLimitInfo { fetch, skip }) - } - - fn extract_limit_from_sort( - sort: &datafusion_expr::logical_plan::Sort, - ) -> Option { - let fetch = sort.fetch?; - Some(VectorLimitInfo { fetch, skip: 0 }) - } - - fn record_vector_hint(&mut self, table_scan: &datafusion_expr::logical_plan::TableScan) { - let Some(limit) = self.current_limit.as_ref() else { - return; - }; - let Some(distance) = self.current_distance.as_ref() else { - return; - }; - // Only emit hints when distance+limit are present and the table matches the sort target. - if let Some(ref hint_table) = distance.table_reference - && hint_table != &table_scan.table_name - { - return; - } - - let non_null_constraint = - self.is_column_non_null(&distance.table_reference, &distance.column_name); - self.vector_hints - .entry(table_scan.table_name.clone()) - .or_default() - .push_back(VectorHintEntry { - distance: distance.clone(), - limit: limit.clone(), - non_null_constraint, - }); - } - - fn take_vector_hint(&mut self, table_name: &TableReference) -> Option { - let hints = self.vector_hints.get_mut(table_name)?; - let hint = hints.pop_front(); - if hints.is_empty() { - self.vector_hints.remove(table_name); - } - hint - } - - fn is_column_non_null( - &self, - table_reference: &Option, - column_name: &str, - ) -> bool { - self.non_null_columns.contains(&ColumnKey { - table_reference: table_reference.clone(), - column_name: column_name.to_string(), - }) - } - - fn extract_query_vector(expr: &Expr) -> Option> { - match expr { - Expr::Literal(scalar, _) => match scalar { - ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => { - Self::parse_json_vector(s) - } - ScalarValue::Binary(Some(bytes)) | ScalarValue::LargeBinary(Some(bytes)) => { - Self::parse_binary_vector(bytes) - } - _ => None, - }, - _ => None, - } - } - - fn parse_json_vector(s: &str) -> Option> { - let trimmed = s.trim(); - if !trimmed.starts_with('[') || !trimmed.ends_with(']') { - return None; - } - parse_string_to_vector_type_value(trimmed, None) - .ok() - .and_then(|bytes| if bytes.is_empty() { None } else { Some(bytes) }) - .and_then(|bytes| Self::parse_binary_vector(&bytes)) - } - - fn parse_binary_vector(bytes: &[u8]) -> Option> { - if bytes.is_empty() || !bytes.len().is_multiple_of(4) { - return None; - } - Some( - bytes - .chunks(4) - .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect(), - ) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use api::v1::SemanticType; - use common_function::function::Function; - use common_function::scalars::udf::create_udf; - use common_function::scalars::vector::distance::{VEC_DOT_PRODUCT, VEC_L2SQ_DISTANCE}; - use datafusion::datasource::DefaultTableSource; - use datafusion::logical_expr::ColumnarValue; - use datafusion_common::{Column, DataFusionError, Result, ScalarValue}; - use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::logical_plan::JoinType; - use datafusion_expr::{ - Expr, Extension, LogicalPlan, LogicalPlanBuilder, Signature, Subquery, Volatility, col, lit, - }; - use datafusion_optimizer::{OptimizerContext, OptimizerRule}; - use datatypes::schema::ColumnSchema; - use promql::extension_plan::InstantManipulate; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; - use store_api::storage::{ConcreteDataType, VectorDistanceMetric, VectorSearchRequest}; - - use super::VectorSearchState; - use crate::dummy_catalog::DummyTableProvider; - use crate::optimizer::scan_hint::ScanHintRule; - use crate::optimizer::test_util::MetaRegionEngine; - - struct TestVectorFunction { - name: &'static str, - signature: Signature, - } - - impl TestVectorFunction { - fn new(name: &'static str) -> Self { - Self { - name, - signature: Signature::any(2, Volatility::Immutable), - } - } - } - - impl std::fmt::Display for TestVectorFunction { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.name) - } - } - - impl Function for TestVectorFunction { - fn name(&self) -> &str { - self.name - } - - fn return_type( - &self, - _input_types: &[datatypes::arrow::datatypes::DataType], - ) -> Result { - Ok(datatypes::arrow::datatypes::DataType::Float32) - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn invoke_with_args( - &self, - _args: datafusion_expr::ScalarFunctionArgs, - ) -> Result { - Err(DataFusionError::Execution( - "test udf should not be invoked".to_string(), - )) - } - } - - fn vec_distance_expr(function_name: &'static str) -> Expr { - vec_distance_expr_with_query(function_name, "[1.0, 2.0]") - } - - fn vec_distance_expr_with_query(function_name: &'static str, query_vector: &str) -> Expr { - let udf = create_udf(Arc::new(TestVectorFunction::new(function_name))); - Expr::ScalarFunction(ScalarFunction::new_udf( - Arc::new(udf), - vec![ - col("v"), - lit(ScalarValue::Utf8(Some(query_vector.to_string()))), - ], - )) - } - - fn scan_request_from_plan(plan: &LogicalPlan) -> store_api::storage::ScanRequest { - let mut request = None; - plan.apply_with_subqueries(|node| { - if let LogicalPlan::TableScan(scan) = node - && let Some(source) = scan.source.downcast_ref::() - && let Some(provider) = source.table_provider.downcast_ref::() - { - request = Some(provider.scan_request()); - } - Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) - }) - .unwrap(); - request.unwrap() - } - - fn vector_hint_from_plan(plan: &LogicalPlan) -> Option { - scan_request_from_plan(plan).vector_search - } - - fn build_dummy_provider(column_id: u32) -> Arc { - build_dummy_provider_with_nullable(column_id, false) - } - - fn build_dummy_provider_with_nullable( - column_id: u32, - nullable_vector: bool, - ) -> Arc { - let mut builder = RegionMetadataBuilder::new(0.into()); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), true), - semantic_type: SemanticType::Tag, - column_id: 1, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 2, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "v", - ConcreteDataType::vector_datatype(2), - nullable_vector, - ), - semantic_type: SemanticType::Field, - column_id, - }) - .primary_key(vec![1]); - let metadata = Arc::new(builder.build().unwrap()); - let engine = Arc::new(MetaRegionEngine::with_metadata(metadata.clone())); - Arc::new(DummyTableProvider::new(0.into(), engine, metadata)) - } - - #[test] - fn test_parse_json_vector() { - assert_eq!( - VectorSearchState::parse_json_vector("[1.0, 2.0, 3.0]"), - Some(vec![1.0, 2.0, 3.0]) - ); - assert_eq!( - VectorSearchState::parse_json_vector("[1.5, -2.3, 0.0]"), - Some(vec![1.5, -2.3, 0.0]) - ); - assert_eq!(VectorSearchState::parse_json_vector("invalid"), None); - assert_eq!(VectorSearchState::parse_json_vector("[]"), None); - assert_eq!(VectorSearchState::parse_json_vector(""), None); - assert_eq!(VectorSearchState::parse_json_vector("["), None); - assert_eq!(VectorSearchState::parse_json_vector("[1.0, abc]"), None); - } - - #[test] - fn test_parse_binary_vector() { - let v1: f32 = 1.0; - let v2: f32 = 2.0; - let mut bytes = Vec::new(); - bytes.extend_from_slice(&v1.to_le_bytes()); - bytes.extend_from_slice(&v2.to_le_bytes()); - - let result = VectorSearchState::parse_binary_vector(&bytes); - assert_eq!(result, Some(vec![1.0, 2.0])); - - let result = VectorSearchState::parse_binary_vector(&[1, 2, 3]); - assert_eq!(result, None); - } - - #[test] - fn test_dummy_provider_vector_hint() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.column_id, 10); - assert_eq!(hint.k, 5); - assert_eq!(hint.metric, VectorDistanceMetric::L2sq); - assert_eq!(hint.query_vector, vec![1.0, 2.0]); - } - - #[test] - fn test_vector_hint_is_applied_when_single_evaluation_limit_blocks_last_row() { - let provider = build_dummy_provider(10); - let source = Arc::new(DefaultTableSource::new(provider)); - let distance = vec_distance_expr(VEC_L2SQ_DISTANCE); - let scan = LogicalPlanBuilder::scan_with_filters("t", source, None, vec![]) - .unwrap() - .sort(vec![distance.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - let plan = LogicalPlan::Extension(Extension { - node: Arc::new(InstantManipulate::new( - 1000, - 1000, - 1000, - 1000, - 0, - "ts".to_string(), - vec![], - Some("v".to_string()), - scan, - )), - }); - - let rewritten = ScanHintRule - .rewrite(plan, &OptimizerContext::default()) - .unwrap() - .data; - let request = scan_request_from_plan(&rewritten); - assert_eq!(request.series_row_selector, None); - assert_eq!( - request.vector_search, - Some(VectorSearchRequest { - column_id: 10, - query_vector: vec![1.0, 2.0], - k: 5, - metric: VectorDistanceMetric::L2sq, - }) - ); - } - - #[test] - fn test_single_evaluation_limit_keeps_vector_hints_on_both_branches() { - let provider_a = build_dummy_provider(10); - let provider_b = build_dummy_provider(20); - let distance_a = vec_distance_expr_with_query(VEC_L2SQ_DISTANCE, "[1.0, 2.0]"); - let distance_b = vec_distance_expr_with_query(VEC_L2SQ_DISTANCE, "[3.0, 4.0]"); - - let branch_a = LogicalPlanBuilder::scan_with_filters( - "t", - Arc::new(DefaultTableSource::new(provider_a)), - None, - vec![], - ) - .unwrap() - .sort(vec![distance_a.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - let branch_a = LogicalPlan::Extension(Extension { - node: Arc::new(InstantManipulate::new( - 1000, - 1000, - 1000, - 1000, - 0, - "ts".to_string(), - vec![], - Some("v".to_string()), - branch_a, - )), - }); - - let branch_b = LogicalPlanBuilder::scan_with_filters( - "t", - Arc::new(DefaultTableSource::new(provider_b)), - None, - vec![], - ) - .unwrap() - .sort(vec![distance_b.sort(true, false)]) - .unwrap() - .limit(0, Some(9)) - .unwrap() - .build() - .unwrap(); - - let plan = LogicalPlanBuilder::from(branch_a) - .union(branch_b) - .unwrap() - .build() - .unwrap(); - let rewritten = ScanHintRule - .rewrite(plan, &OptimizerContext::default()) - .unwrap() - .data; - - let LogicalPlan::Union(union) = rewritten else { - panic!("expected union plan") - }; - let request_a = scan_request_from_plan(&union.inputs[0]); - let request_b = scan_request_from_plan(&union.inputs[1]); - - assert_eq!(request_a.series_row_selector, None); - assert_eq!( - request_a.vector_search, - Some(VectorSearchRequest { - column_id: 10, - query_vector: vec![1.0, 2.0], - k: 5, - metric: VectorDistanceMetric::L2sq, - }) - ); - assert_eq!(request_b.series_row_selector, None); - assert_eq!( - request_b.vector_search, - Some(VectorSearchRequest { - column_id: 20, - query_vector: vec![3.0, 4.0], - k: 9, - metric: VectorDistanceMetric::L2sq, - }) - ); - - // The catalog providers are not mutated; the requests above belong to - // the two rewritten scan use-sites. - assert_ne!(request_a.vector_search, request_b.vector_search); - } - - #[test] - fn test_limit_offset_for_vector_hint() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(5, Some(10)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.k, 15); - } - - #[test] - fn test_inner_product_sort_direction() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_DOT_PRODUCT); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source.clone(), None, vec![]) - .unwrap() - .sort(vec![expr.clone().sort(true, false)]) - .unwrap() - .limit(0, Some(3)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - assert!(dummy_provider.get_vector_search_hint().is_none()); - - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(false, false)]) - .unwrap() - .limit(0, Some(3)) - .unwrap() - .build() - .unwrap(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.metric, VectorDistanceMetric::InnerProduct); - } - - #[test] - fn test_no_limit_clause() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(dummy_provider.get_vector_search_hint().is_none()); - } - - #[test] - fn test_nullable_vector_requires_is_not_null_filter() { - let dummy_provider = build_dummy_provider_with_nullable(10, true); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(dummy_provider.get_vector_search_hint().is_none()); - } - - #[test] - fn test_nullable_vector_with_is_not_null_filter() { - let dummy_provider = build_dummy_provider_with_nullable(10, true); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .filter(col("v").is_not_null()) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.column_id, 10); - assert_eq!(hint.k, 5); - } - - #[test] - fn test_sort_fetch_for_vector_hint() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort_with_limit(vec![expr.sort(true, false)], Some(4)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.k, 4); - } - - #[test] - fn test_limit_scoped_to_sort_branch() { - let t1_provider = build_dummy_provider(10); - let t2_provider = build_dummy_provider(20); - let t1_source = Arc::new(DefaultTableSource::new(t1_provider.clone())); - let t2_source = Arc::new(DefaultTableSource::new(t2_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - - let left = LogicalPlanBuilder::scan_with_filters("t1", t1_source, None, vec![]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let right = LogicalPlanBuilder::scan_with_filters("t2", t2_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .build() - .unwrap(); - - let join_keys: (Vec, Vec) = (vec![], vec![]); - let plan = LogicalPlanBuilder::from(left) - .join(right, JoinType::Inner, join_keys, None) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(t1_provider.get_vector_search_hint().is_none()); - assert!(t2_provider.get_vector_search_hint().is_none()); - } - - fn vec_distance_expr_qualified( - function_name: &'static str, - table_name: &str, - column_name: &str, - ) -> Expr { - use datafusion_common::Column; - - let udf = create_udf(Arc::new(TestVectorFunction::new(function_name))); - let qualified_col = Expr::Column(Column::new(Some(table_name.to_string()), column_name)); - Expr::ScalarFunction(ScalarFunction::new_udf( - Arc::new(udf), - vec![ - qualified_col, - lit(ScalarValue::Utf8(Some("[1.0, 2.0]".to_string()))), - ], - )) - } - - #[test] - fn test_no_vector_hint_above_join() { - let t1_provider = build_dummy_provider(10); - let t2_provider = build_dummy_provider(20); - let t1_source = Arc::new(DefaultTableSource::new(t1_provider.clone())); - let t2_source = Arc::new(DefaultTableSource::new(t2_provider.clone())); - - let left = LogicalPlanBuilder::scan_with_filters("t1", t1_source, None, vec![]) - .unwrap() - .build() - .unwrap(); - - let right = LogicalPlanBuilder::scan_with_filters("t2", t2_source, None, vec![]) - .unwrap() - .build() - .unwrap(); - - let join_keys: (Vec, Vec) = (vec![], vec![]); - let join_plan = LogicalPlanBuilder::from(left) - .join(right, JoinType::Inner, join_keys, None) - .unwrap() - .build() - .unwrap(); - - let expr = vec_distance_expr_qualified(VEC_L2SQ_DISTANCE, "t1", "v"); - let plan = LogicalPlanBuilder::from(join_plan) - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(t1_provider.get_vector_search_hint().is_none()); - assert!(t2_provider.get_vector_search_hint().is_none()); - } - - // Simple subqueries (without non-inlineable ops like Limit/Sort/Distinct/Aggregate/Window) - // allow hints to propagate through. See known limitations in VectorSearchState docs. - #[test] - fn test_simple_subquery_allows_hint_propagation() { - let provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(provider.clone())); - let scan_plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .build() - .unwrap(); - - let subquery = LogicalPlan::Subquery(Subquery { - subquery: Arc::new(scan_plan), - outer_ref_columns: vec![], - spans: Default::default(), - }); - - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::from(subquery) - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - // Hint propagates through simple subquery - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.k, 5); - } - - // Subqueries with non-inlineable ops (Limit/Sort/Distinct/Aggregate/Window) block hint propagation. - #[test] - fn test_subquery_with_limit_blocks_hint() { - let provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(provider.clone())); - let scan_plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .limit(0, Some(100)) // non-inlineable op inside subquery - .unwrap() - .build() - .unwrap(); - - let subquery = LogicalPlan::Subquery(Subquery { - subquery: Arc::new(scan_plan), - outer_ref_columns: vec![], - spans: Default::default(), - }); - - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::from(subquery) - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - // Hint does NOT propagate through subquery with non-inlineable ops - assert!(provider.get_vector_search_hint().is_none()); - } - - #[test] - fn test_qualified_column_scopes_hint_to_correct_table() { - let t1_provider = build_dummy_provider(10); - let t2_provider = build_dummy_provider(20); - let t1_source = Arc::new(DefaultTableSource::new(t1_provider.clone())); - let t2_source = Arc::new(DefaultTableSource::new(t2_provider.clone())); - - let expr = vec_distance_expr_qualified(VEC_L2SQ_DISTANCE, "t2", "v"); - - let t1_plan = LogicalPlanBuilder::scan_with_filters("t1", t1_source, None, vec![]) - .unwrap() - .sort(vec![expr.clone().sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let t2_plan = LogicalPlanBuilder::scan_with_filters("t2", t2_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - - let _ = ScanHintRule.rewrite(t1_plan, &context).unwrap(); - assert!(t1_provider.get_vector_search_hint().is_none()); - - let rewritten = ScanHintRule.rewrite(t2_plan, &context).unwrap().data; - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.column_id, 20); - assert_eq!(hint.k, 5); - } -} diff --git a/src/query/src/optimizer/windowed_sort.rs b/src/query/src/optimizer/windowed_sort.rs index 3e4ccaf91c4..800bc103d38 100644 --- a/src/query/src/optimizer/windowed_sort.rs +++ b/src/query/src/optimizer/windowed_sort.rs @@ -199,12 +199,12 @@ fn fetch_partition_range(input: Arc) -> DataFusionResult StatusCode::InvalidArguments, diff --git a/src/query/src/promql/label_values.rs b/src/query/src/promql/label_values.rs index 44c173477bb..cae53c64853 100644 --- a/src/query/src/promql/label_values.rs +++ b/src/query/src/promql/label_values.rs @@ -33,19 +33,26 @@ use crate::promql::error::{ /// a negative duration, and an RFC3339 query parameter can name one, so the sign /// is recovered here instead of unwrapping. fn millis_since_epoch(time: SystemTime) -> Result { + Ok(Timestamp::new_millisecond(signed_millis_since_epoch(time)?)) +} + +/// Converts a [`SystemTime`] to signed milliseconds since the Unix epoch, rejecting instants +/// outside the `i64` millisecond range. +/// +/// Sub-millisecond parts round down, also before the epoch, like Prometheus' +/// `timestamp.FromTime`. +pub(crate) fn signed_millis_since_epoch(time: SystemTime) -> Result { let (millis, before_epoch) = match time.duration_since(UNIX_EPOCH) { Ok(duration) => (duration.as_millis(), false), - Err(earlier) => (earlier.duration().as_millis(), true), + Err(earlier) => (earlier.duration().as_nanos().div_ceil(1_000_000), true), }; - let millis = i64::try_from(millis) - .ok() - .with_context(|| SystemTimeOutOfRangeSnafu { time })?; + signed_millis(millis, before_epoch).with_context(|| SystemTimeOutOfRangeSnafu { time }) +} - Ok(Timestamp::new_millisecond(if before_epoch { - -millis - } else { - millis - })) +/// Converts a millisecond magnitude and epoch direction without platform time limits. +fn signed_millis(millis: u128, before_epoch: bool) -> Option { + let millis = i64::try_from(millis).ok()?; + Some(if before_epoch { -millis } else { millis }) } fn build_time_filter(time_index_expr: Expr, start: Timestamp, end: Timestamp) -> Expr { @@ -139,7 +146,16 @@ mod tests { #[test] fn millis_beyond_i64_are_rejected() { - let time = UNIX_EPOCH + Duration::from_secs(1 << 60); - assert!(millis_since_epoch(time).is_err()); + let limit = i64::MAX as u128; + assert_eq!(signed_millis(limit, false), Some(i64::MAX)); + assert_eq!(signed_millis(limit, true), Some(-i64::MAX)); + assert_eq!(signed_millis(limit + 1, false), None); + assert_eq!(signed_millis(limit + 1, true), None); + + // Windows cannot represent this SystemTime. The conversion boundary is + // tested above on every platform; exercise the wrapper where possible. + if let Some(time) = UNIX_EPOCH.checked_add(Duration::from_millis(i64::MAX as u64 + 1)) { + assert!(millis_since_epoch(time).is_err()); + } } } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 94081c2d284..5867301f972 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -12,11 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod at_modifier; +mod function_plans; +mod island; mod matching_filters; -use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; +mod set_operator; + +use std::collections::{BTreeSet, HashSet, VecDeque}; use std::sync::Arc; -use std::time::UNIX_EPOCH; use arrow::datatypes::IntervalDayTime; use async_recursion::async_recursion; @@ -43,55 +47,48 @@ use datafusion::functions_window::row_number::RowNumber; use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction}; use datafusion::logical_expr::expr_rewriter::normalize_cols; use datafusion::logical_expr::{ - BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator, + BinaryExpr, Cast, EmptyRelation, Extension, LogicalPlan, LogicalPlanBuilder, Operator, ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition, }; -use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::prelude as df_prelude; use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; use datafusion::scalar::ScalarValue; -use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchema, NullEquality, TableReference}; use datafusion_expr::expr::WindowFunctionParams; use datafusion_expr::expr_fn::when; -use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{conjunction, disjunction}; -use datafusion_expr::{ - ExprSchemable, Literal, Projection, SortExpr, TableScanBuilder, TableSource, col, lit, -}; +use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, ident, lit}; use datafusion_functions::core::coalesce; use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit}; use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; use itertools::Itertools; use once_cell::sync::Lazy; use promql::extension_plan::{ - Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond, - RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, + EmptyMetric, InstantManipulate, Millisecond, RangeManipulate, SeriesDivide, SeriesNormalize, build_special_time_expr, }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, + IDelta, Increase, LastOverTime, MatchGroupViolation, MaxOverTime, MinOverTime, MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, - NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, - NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, - NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, - NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, - NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, - NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, - Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramIDelta, + NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, + NativeHistogramPresentOverTime, NativeHistogramRate, NativeHistogramResets, + NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, + NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, PredictLinear, + PresentOverTime, PromqlFloatToString, QuantileOverTime, Rate, Resets, Round, StddevOverTime, + StdvarOverTime, SumOverTime, UniqueMatchGroup, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr, - Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral, - Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality, - VectorSelector, token, + Function, LabelModifier, MatrixSelector, NumberLiteral, Offset, ParenExpr, StringLiteral, + SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector, token, }; use regex::{self, Regex}; use snafu::{OptionExt, ResultExt, ensure}; @@ -106,17 +103,20 @@ use crate::parser::{ EXPLAIN_VERBOSE_NODE_NAME, }; use crate::promql::error::{ - CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu, - ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, - InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, - MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu, - Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, + CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectRangeSelectorSnafu, + FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, InvalidRegularExpressionSnafu, + InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, MultipleMetricMatchersSnafu, + MultipleVectorSnafu, NoMetricMatcherSnafu, RangeVectorInRangeQuerySnafu, Result, + SubqueryTimestampOutOfRangeSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu, - UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu, - ZeroRangeSelectorSnafu, + UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, ZeroSubqueryStepSnafu, }; +use crate::promql::label_values::signed_millis_since_epoch; use crate::query_engine::QueryEngineState; +/// Step of a subquery without an explicit one (`x[5m:]`). Prometheus uses the global +/// evaluation interval, whose default is one minute, rather than the query step. +const DEFAULT_SUBQUERY_STEP_MS: Millisecond = 60_000; /// `time()` function in PromQL. const SPECIAL_TIME_FUNCTION: &str = "time"; /// `scalar()` function in PromQL. @@ -149,11 +149,11 @@ const FIELD_COLUMN_MATCHER: &str = "__field__"; const SCHEMA_COLUMN_MATCHER: &str = "__schema__"; const DB_COLUMN_MATCHER: &str = "__database__"; -/// Prefix for generated binary island leaf aliases. -const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_"; const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_"; const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_"; +/// Per-match-group row count the vector matching cardinality check is built on. +const MATCH_GROUP_COUNT_COLUMN: &str = "__promql_match_group_count"; /// Threshold for scatter scan mode const MAX_SCATTER_POINTS: i64 = 400; @@ -168,6 +168,12 @@ struct PromPlannerContext { end: Millisecond, interval: Millisecond, lookback_delta: Millisecond, + /// Evaluation range of the whole statement, which `@ start()` and `@ end()` refer to. + /// + /// Unlike [`Self::start`] and [`Self::end`], these are never rewritten while planning, so a + /// selector inside a subquery still resolves `@ start()` / `@ end()` against the statement. + stmt_start: Millisecond, + stmt_end: Millisecond, // planner states table_name: Option, @@ -196,189 +202,61 @@ struct PromPlannerContext { schema_name: Option, /// The range in millisecond of range selector. None if there is no range selector. range: Option, + /// The offset in milliseconds the window of the last planned range selector is folded with, + /// or `None` when no range selector has been planned since the last read. + /// + /// [`Self::start`] and the sample timestamps are compared on the shifted timeline: a range + /// payload carries `sample_timestamp + offset`, while the time index column of a folded row + /// stays the evaluation timestamp of its step. A function that reads it right after its input + /// plan is built, like `predict_linear`, consumes it instead of reading the state, so the + /// offset cannot leak from one input to another; see + /// [`PromPlanner::create_function_expr`]. + range_fold_offset: Option, } -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -struct VectorLeafKey { - metric_name: String, - matchers: Vec<(String, String, String)>, - or_matchers: Vec>, - offset_ms: i128, - at: String, -} - -#[derive(Debug, Clone)] -struct IslandLeaf { - selector: VectorSelector, - display_table: String, -} - -#[derive(Debug, Clone)] -enum IslandExpr { - VectorLeaf(usize), - Scalar(DfExpr), - Unary { - input: Box, - }, - Binary { - op: TokenType, - lhs: Box, - rhs: Box, - }, -} - -impl IslandExpr { - fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option { - if let Some(expr) = PromPlanner::try_build_literal_expr(expr) { - return Some(Self::Scalar(expr)); - } - - match expr { - PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env), - PromExpr::VectorSelector(selector) => { - let leaf = env.intern_leaf(selector)?; - Some(Self::VectorLeaf(leaf)) - } - PromExpr::Unary(UnaryExpr { expr }) => { - let input = Self::try_new(expr, env)?; - Some(Self::Unary { - input: Box::new(input), - }) - } - PromExpr::Binary(PromBinaryExpr { - lhs, - rhs, - op, - modifier, - }) if matches!( - op.id(), - token::T_ADD - | token::T_SUB - | token::T_MUL - | token::T_DIV - | token::T_MOD - | token::T_POW - | token::T_ATAN2 - ) && modifier.as_ref().is_none_or(|modifier| { - !modifier.return_bool - && modifier.matching.is_none() - && matches!(modifier.card, VectorMatchCardinality::OneToOne) - && modifier.fill_values.lhs.is_none() - && modifier.fill_values.rhs.is_none() - }) => - { - let lhs = Self::try_new(lhs, env)?; - let rhs = Self::try_new(rhs, env)?; - Some(Self::Binary { - op: *op, - lhs: Box::new(lhs), - rhs: Box::new(rhs), - }) - } - _ => None, - } - } -} - -#[derive(Debug, Default)] -struct IslandCollectEnv { - leaf_by_key: HashMap, - leaves: Vec, - vector_occurrences: usize, -} - +/// Result labels a vector-vector binary operation derives from its matching modifier, projected +/// from the operand each label belongs to. #[derive(Debug)] -struct PlannedIslandLeaf { - plan: LogicalPlan, - ctx: PromPlannerContext, - alias: TableReference, - display_table: String, -} - -#[derive(Debug)] -struct IslandFieldExprs { +struct BinaryResultLabels { exprs: Vec, names: Vec, - scalar: bool, + aggregation_field_labels: Vec, + /// `__tsid` column of the operand the labels come from, when that operand contributes its + /// whole tag set and the column still identifies the result series. + tsid: Option, } -impl VectorLeafKey { - fn from_selector(selector: &VectorSelector) -> Option { - let mut metric_name = selector.name.clone(); - let mut matchers = Vec::with_capacity(selector.matchers.matchers.len()); - let matcher_key = |matcher: &Matcher| { - ( - matcher.name.clone(), - matcher.op.to_string(), - matcher.value.clone(), - ) - }; - - for matcher in &selector.matchers.matchers { - if matcher.name == METRIC_NAME { - if matcher.op != MatchOp::Equal || metric_name.is_some() { - return None; - } - metric_name = Some(matcher.value.clone()); - } else { - matchers.push(matcher_key(matcher)); - } - } - matchers.sort(); - - let mut or_matchers = selector - .matchers - .or_matchers - .iter() - .map(|group| { - let mut group = group.iter().map(matcher_key).collect::>(); - group.sort(); - group - }) - .collect::>(); - or_matchers.sort(); - - Some(Self { - metric_name: metric_name?, - matchers, - or_matchers, - offset_ms: match &selector.offset { - Some(Offset::Pos(duration)) => duration.as_millis() as i128, - Some(Offset::Neg(duration)) => -(duration.as_millis() as i128), - None => 0, - }, - at: format!("{:?}", selector.at), - }) +impl BinaryResultLabels { + fn apply(&self, ctx: &mut PromPlannerContext) { + ctx.tag_columns = self.names.clone(); + ctx.aggregation_field_labels = self.aggregation_field_labels.clone(); + ctx.use_tsid = self.tsid.is_some(); } -} -impl IslandCollectEnv { - fn intern_leaf(&mut self, selector: &VectorSelector) -> Option { - self.vector_occurrences += 1; - let key = VectorLeafKey::from_selector(selector)?; - if let Some(id) = self.leaf_by_key.get(&key) { - return Some(*id); - } - - let id = self.leaves.len(); - self.leaves.push(IslandLeaf { - selector: selector.clone(), - display_table: key.metric_name.clone(), - }); - self.leaf_by_key.insert(key, id); - Some(id) + /// `__tsid` is projected from whichever operand the labels come from, so it carries that + /// operand's qualifier. Re-qualify it as the result's own, which is what the context names + /// and what the enclosing expression looks the column up by. + fn tsid_projection(&self, table_ref: Option) -> Option { + self.tsid + .clone() + .map(|tsid| tsid.alias_qualified(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)) } } impl PromPlannerContext { - fn from_eval_stmt(stmt: &EvalStmt) -> Self { - Self { - start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _, - end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _, + fn from_eval_stmt(stmt: &EvalStmt) -> Result { + // Evaluation times before the Unix epoch are valid in Prometheus. + let start = signed_millis_since_epoch(stmt.start)?; + let end = signed_millis_since_epoch(stmt.end)?; + Ok(Self { + start, + end, interval: stmt.interval.as_millis() as _, lookback_delta: stmt.lookback_delta.as_millis() as _, + stmt_start: start, + stmt_end: end, ..Default::default() - } + }) } /// Reset all planner states @@ -392,6 +270,7 @@ impl PromPlannerContext { self.selector_matcher.clear(); self.schema_name = None; self.range = None; + self.range_fold_offset = None; } /// Reset table name and schema to empty @@ -434,12 +313,12 @@ impl PromPlanner { ) -> Result { let mut planner = Self { table_provider, - ctx: PromPlannerContext::from_eval_stmt(stmt), + ctx: PromPlannerContext::from_eval_stmt(stmt)?, promql_annotations, }; let plan = planner - .prom_expr_to_plan(&stmt.expr, query_engine_state) + .prom_root_expr_to_plan(&stmt.expr, query_engine_state) .await?; // Never leak internal series identifier to output. @@ -471,6 +350,16 @@ impl PromPlanner { timestamp_fn: bool, query_engine_state: &QueryEngineState, ) -> Result { + // An anchored range call is step-invariant: evaluate it once, at the start of the + // evaluation, and report its result at every step; see + // [`Self::promote_anchored_range_call`]. + if let Some(plan) = self + .promote_anchored_range_call(prom_expr, timestamp_fn, query_engine_state) + .await? + { + return Ok(plan); + } + let res = match prom_expr { PromExpr::Aggregate(expr) => { self.prom_aggr_expr_to_plan(query_engine_state, expr) @@ -513,24 +402,230 @@ impl PromPlanner { Ok(res) } + /// Plans the root expression of a statement, or the expression wrapped by a root extension + /// such as `EXPLAIN` or an alias. + /// + /// Prometheus only accepts a range-vector expression in an instant query, whose result is + /// the matrix of samples the expression selects. + async fn prom_root_expr_to_plan( + &mut self, + prom_expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result { + if matches!(prom_expr, PromExpr::Extension(_)) + || prom_expr.value_type() != ValueType::Matrix + { + return self.prom_expr_to_plan(prom_expr, query_engine_state).await; + } + ensure!(self.ctx.start == self.ctx.end, RangeVectorInRangeQuerySnafu); + self.prom_range_vector_result_to_plan(prom_expr, query_engine_state) + .await + } + + /// Plans the result of a range-vector expression in an instant query: the raw samples of a + /// range selector, or the evaluation points of a subquery, with their own timestamps. + /// + /// Range functions consume the same expressions through [`RangeManipulate`], whose range + /// arrays are an internal representation that must not reach the query output. + async fn prom_range_vector_result_to_plan( + &mut self, + prom_expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result { + let mut expr = prom_expr; + while let PromExpr::Paren(ParenExpr { expr: inner }) = expr { + expr = inner; + } + match expr { + PromExpr::MatrixSelector(selector) => { + self.prom_range_selector_samples_to_plan(selector).await + } + PromExpr::Subquery(subquery) => { + self.prom_subquery_samples_to_plan(subquery, query_engine_state) + .await + } + _ => UnsupportedExprSnafu { + name: format!("range vector result of {expr:?}"), + } + .fail(), + } + } + + /// Selects the samples of `selector` in `(t - offset - range, t - offset]`, where `t` is the + /// evaluation time or the `@` anchor. Stale markers are dropped and timestamps are kept. + async fn prom_range_selector_samples_to_plan( + &mut self, + selector: &MatrixSelector, + ) -> Result { + let MatrixSelector { vs, range } = selector; + let VectorSelector { + name, + offset, + matchers, + at, + } = vs; + let matchers = self.preprocess_label_matchers(matchers, name)?; + ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); + self.ctx.range = Some(range.as_millis() as _); + let offset_ms = match self.at_modifier_offset(at, offset)? { + Some(at_offset) => at_offset, + None => Self::offset_millis(offset), + }; + if let Some(empty_plan) = self.setup_context().await? { + return Ok(empty_plan); + } + let plan = self + .selector_to_series_normalize_plan(offset_ms, matchers, true) + .await?; + self.millisecond_time_index_plan(plan) + } + + /// Casts a microsecond or nanosecond time index to milliseconds, the precision of every + /// other PromQL result and the only one the Prometheus HTTP API encodes. + fn millisecond_time_index_plan(&self, plan: LogicalPlan) -> Result { + let time_index = + self.ctx + .time_index_column + .as_deref() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let (_, field) = plan + .schema() + .qualified_field_with_unqualified_name(time_index) + .context(DataFusionPlanningSnafu)?; + let millis = match field.data_type() { + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => return Ok(plan), + ArrowDataType::Timestamp(_, timezone) => { + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, timezone.clone()) + } + other => unreachable!("time index is a timestamp, got {other}"), + }; + let exprs = plan + .schema() + .iter() + .map(|(qualifier, field)| { + let column = DfExpr::Column(Column::new(qualifier.cloned(), field.name())); + if field.name() == time_index { + column + .cast_to(&millis, plan.schema()) + .map(|cast| cast.alias_qualified(qualifier.cloned(), field.name().clone())) + } else { + Ok(column) + } + }) + .collect::>>() + .context(DataFusionPlanningSnafu)?; + LogicalPlanBuilder::from(plan) + .project(exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + /// Evaluates the inner expression of `subquery` at the multiples of its step in + /// `(t - offset - range, t - offset]`, where `t` is the evaluation time or the `@` anchor. + async fn prom_subquery_samples_to_plan( + &mut self, + subquery: &SubqueryExpr, + query_engine_state: &QueryEngineState, + ) -> Result { + let SubqueryExpr { + expr, + range, + step, + offset, + at, + .. + } = subquery; + ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); + let end = match self.at_ref_time(at, offset)? { + Some(anchor) => anchor, + None => Self::anchor_sub(self.ctx.end, Self::offset_millis(offset))?, + }; + self.ctx.interval = match step { + Some(step) => step.as_millis() as _, + None => DEFAULT_SUBQUERY_STEP_MS, + }; + let interval = self.ctx.interval; + ensure!(interval > 0, ZeroSubqueryStepSnafu); + let window = Self::subquery_child_window(end, end, 0, 0, range.as_millis() as _, interval)?; + let Some((start, end)) = window else { + // No step multiple falls in the window; plan the inner expression only for its + // schema. + self.ctx.start = end; + self.ctx.end = end; + let plan = self.prom_expr_to_plan(expr, query_engine_state).await?; + return Ok(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: plan.schema().clone(), + })); + }; + self.ctx.start = start; + self.ctx.end = end; + self.prom_expr_to_plan(expr, query_engine_state).await + } + async fn prom_subquery_expr_to_plan( &mut self, query_engine_state: &QueryEngineState, subquery_expr: &SubqueryExpr, ) -> Result { let SubqueryExpr { - expr, range, step, .. + expr, + range, + step, + offset, + .. } = subquery_expr; + // Shift the child window back; `RangeManipulate` restores the outer timeline. + 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 current_interval = self.ctx.interval; if let Some(step) = step { self.ctx.interval = step.as_millis() as _; } + ensure!(self.ctx.interval > 0, ZeroRangeSelectorSnafu); let current_start = self.ctx.start; - self.ctx.start -= range.as_millis() as i64 - self.ctx.interval; + let current_end = self.ctx.end; + // Reject an invalid parent window rather than treating it as empty. + ensure!( + current_start <= current_end, + InvalidTimeRangeSnafu { + start: current_start, + end: current_end, + } + ); + // An empty child needs only its schema; plan it over the valid caller window. + let child_window = Self::subquery_child_window( + current_start, + current_end, + current_interval, + offset_ms, + range.as_millis() as Millisecond, + self.ctx.interval, + )?; + if let Some((child_start, child_end)) = child_window { + self.ctx.start = child_start; + self.ctx.end = child_end; + } let input = self.prom_expr_to_plan(expr, query_engine_state).await?; self.ctx.interval = current_interval; self.ctx.start = current_start; + self.ctx.end = current_end; + // Return no rows with the child schema so enclosing fallbacks still work. + let input = match child_window { + Some(_) => input, + None => LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: input.schema().clone(), + }), + }; ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); let range_ms = range.as_millis() as _; @@ -587,23 +682,42 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)?; let divide_plan = LogicalPlan::Extension(Extension { node: Arc::new(SeriesDivide::new( - series_key_columns, + series_key_columns.clone(), time_index_column.clone(), sort_plan, )), }); + // `RangeManipulate` has no offset in its protobuf message; decoding recovers it from the + // `SeriesNormalize` directly below. Stale markers are not filtered: the input is computed. + let divide_plan = if offset_ms != 0 { + LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + offset_ms, + time_index_column.clone(), + false, + series_key_columns, + divide_plan, + )), + }) + } else { + divide_plan + }; + let manipulate = RangeManipulate::new( self.ctx.start, self.ctx.end, self.ctx.interval, - 0, + offset_ms, range_ms, time_index_column, self.ctx.field_columns.clone(), divide_plan, ) .context(DataFusionPlanningSnafu)?; + // The payload timestamps are shifted by the subquery offset; see + // [`Self::create_range_eval_ts_expr`]. + self.ctx.range_fold_offset = Some(offset_ms); Ok(LogicalPlan::Extension(Extension { node: Arc::new(manipulate), @@ -622,40 +736,12 @@ impl PromPlanner { param, } = aggr_expr; - let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?; + let input = self.prom_expr_to_plan(expr, query_engine_state).await?; let input_has_tsid = input.schema().fields().iter().any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME && field.data_type() == &ArrowDataType::UInt64 }); - // `__tsid` based scan projection may prune tag columns. Ensure tags referenced in - // aggregation modifiers (`by`/`without`) are available before planning group keys. - let required_group_tags = match modifier { - None => BTreeSet::new(), - Some(LabelModifier::Include(labels)) => labels - .labels - .iter() - .filter(|label| !is_metric_engine_internal_column(label.as_str())) - .cloned() - .collect(), - Some(LabelModifier::Exclude(labels)) => { - let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?; - for label in &labels.labels { - let _ = all_tags.remove(label); - } - all_tags - } - }; - - if !required_group_tags.is_empty() - && required_group_tags - .iter() - .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none()) - { - input = self.ensure_tag_columns_available(input, &required_group_tags)?; - self.refresh_tag_columns_from_schema(input.schema()); - } - match (*op).id() { token::T_TOPK | token::T_BOTTOMK => { self.prom_topk_bottomk_to_plan(aggr_expr, input).await @@ -712,19 +798,82 @@ impl PromPlanner { let label = Self::get_param_value_as_str(*op, param)?; // `count_values` must be grouped by fields, // and project the fields to the new label. - let count_value_exprs = prev_field_exprs.iter().map(|expr| { - match expr { - DfExpr::Column(column) => DfExpr::Column(column.clone()), - _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())), - } - .alias(label) - }); + // + // The generated label is a real label (column) of the output, so it must be + // registered in `ctx.tag_columns` below. Otherwise enclosing expressions + // rebuild their projections from `ctx.tag_columns` and silently drop it. + // + // PromQL sets the generated label *before* the grouping key is built, so an + // input label with the same name is overwritten by the sample value and must + // not remain a grouping key either: samples are grouped by the generated + // label only. Dropping it from the projected tag columns is required as well: + // projecting both would emit two columns with the same name (rejected as an + // ambiguous reference). + self.ctx.tag_columns.retain(|tag| tag != label); + group_exprs.retain( + |expr| !matches!(expr, DfExpr::Column(column) if column.name == label), + ); + // The tag columns projected below are unqualified `Column` references, so they + // end up qualified with whatever qualifier they carry in the input plan. Give + // the generated label the same qualifier, otherwise qualified references to it + // (e.g. from an enclosing binary expression or a vector join) fail to resolve. + let label_qualifier = self + .ctx + .time_index_column + .as_deref() + .and_then(|time_index| { + builder + .schema() + .qualified_field_with_unqualified_name(time_index) + .ok() + }) + .and_then(|(qualifier, _)| qualifier.cloned()); + // The generated label carries the *sample value*, so it must be materialized + // as a string using Prometheus' format + // (`strconv.FormatFloat(value, 'f', -1, 64)`: shortest decimal form without + // an exponent, `1.0` becomes "1"). Labels are inferred from string columns + // when a result is converted into the Prometheus HTTP API JSON format, so a + // numeric label column would be mistaken for the sample value of the series. + let count_value_exprs = prev_field_exprs + .iter() + .map(|expr| { + let value = match expr { + DfExpr::Column(column) => { + let value = DfExpr::Column(column.clone()); + // `prom_float_to_string` formats exactly like Prometheus, + // while arrow's `Float64 -> Utf8` cast would render `1.0`. + let value = if matches!( + builder.schema().field_with_unqualified_name(&column.name), + Ok(field) if field.data_type() == &ArrowDataType::Float64 + ) { + value + } else { + DfExpr::Cast(Cast::new( + Box::new(value), + ArrowDataType::Float64, + )) + }; + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![value], + }) + } + // The value is already formatted (e.g. a native histogram is + // converted to its string form by the aggregation), and the + // aggregate output names it by its schema name. + _ => DfExpr::Column(Column::from_name( + expr.schema_name().to_string(), + )), + }; + DfExpr::Alias(Alias::new(value, label_qualifier.clone(), label)) + }) + .collect::>(); let aggregate_group_exprs = group_exprs .iter() .cloned() .chain(prev_field_exprs.clone()) .collect::>(); - group_exprs.push(col(label)); + group_exprs.push(ident(label)); let project_fields = self .create_field_column_exprs()? .into_iter() @@ -732,11 +881,14 @@ impl PromPlanner { .chain(Some(self.create_time_index_column_expr()?)) .chain(count_value_exprs); - builder + let builder = builder .aggregate(aggregate_group_exprs, aggr_exprs) .context(DataFusionPlanningSnafu)? .project(project_fields) - .context(DataFusionPlanningSnafu)? + .context(DataFusionPlanningSnafu)?; + // The label only exists in the output schema from here on. + self.ctx.tag_columns.push(label.to_string()); + builder } else { builder .aggregate(group_exprs.clone(), aggr_exprs) @@ -841,10 +993,10 @@ impl PromPlanner { ), self.promql_annotations.clone(), )), - args: vec![col(&histogram_column)], + args: vec![ident(&histogram_column)], }); - let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram) - .otherwise(col(&float_column).is_not_null()) + let keep_float = when(ident(&histogram_column).is_not_null(), drop_histogram) + .otherwise(ident(&float_column).is_not_null()) .context(DataFusionPlanningSnafu)?; input = LogicalPlanBuilder::from(input) .filter(keep_float) @@ -895,7 +1047,7 @@ impl PromPlanner { .iter() .fold(None, |expr, rank| { let predicate = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(col(rank)), + left: Box::new(ident(rank)), op: Operator::LtEq, right: Box::new(val.clone()), }); @@ -911,7 +1063,7 @@ impl PromPlanner { }) .unwrap(); - let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect(); + let rank_columns: Vec<_> = rank_columns.into_iter().map(ident).collect(); let mut new_group_exprs = group_exprs.clone(); // Order by ranks @@ -964,338 +1116,16 @@ impl PromPlanner { if Self::field_column_is_native_histogram(&input_schema, col) { Ok(DfExpr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramNeg::scalar_udf()), - args: vec![DfExpr::Column(col.into())], + args: vec![DfExpr::Column(Column::from_name(col))], })) } else { - Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into())))) + Ok(DfExpr::Negative(Box::new(DfExpr::Column( + Column::from_name(col), + )))) } }) } - async fn try_plan_binary_island( - &mut self, - binary_expr: &PromBinaryExpr, - ) -> Result> { - let original_ctx = self.ctx.clone(); - let mut collect_env = IslandCollectEnv::default(); - let Some(island_expr) = - IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env) - else { - return Ok(None); - }; - - if collect_env.leaves.is_empty() - || collect_env.vector_occurrences <= collect_env.leaves.len() - { - return Ok(None); - } - - let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len()); - for (idx, leaf) in collect_env.leaves.iter().enumerate() { - let plan = self - .prom_vector_selector_to_plan(&leaf.selector, false) - .await?; - let ctx = self.ctx.clone(); - let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}")); - let plan = LogicalPlanBuilder::from(plan) - .alias(alias.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - planned_leaves.push(PlannedIslandLeaf { - plan, - ctx, - alias, - display_table: leaf.display_table.clone(), - }); - } - - if planned_leaves.iter().any(|leaf| { - Self::field_columns_contain_native_histogram( - leaf.plan.schema(), - &leaf.ctx.field_columns, - ) - }) { - self.ctx = original_ctx; - return Ok(None); - } - - if !Self::binary_island_join_contexts_supported(&planned_leaves) { - self.ctx = original_ctx; - return Ok(None); - } - - let mut input = planned_leaves[0].plan.clone(); - for right_idx in 1..planned_leaves.len() { - input = self.join_binary_island_leaf( - input, - &planned_leaves[0], - &planned_leaves[right_idx], - )?; - } - - let field_exprs = - Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?; - if field_exprs.scalar || field_exprs.exprs.is_empty() { - self.ctx = original_ctx; - return Ok(None); - } - - let plan = self.project_binary_island( - input, - &planned_leaves[0].alias, - &planned_leaves[0].ctx, - field_exprs, - )?; - Ok(Some(plan)) - } - - fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool { - if leaves - .iter() - .any(|leaf| leaf.ctx.time_index_column.is_none()) - { - return false; - } - - if leaves.len() <= 1 { - return true; - } - - let first_tags = leaves[0].ctx.tag_columns.iter().collect::>(); - - leaves.iter().skip(1).all(|leaf| { - (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan)) - || leaf.ctx.tag_columns.iter().collect::>() == first_tags - }) - } - - fn join_binary_island_leaf( - &self, - left: LogicalPlan, - first_leaf: &PlannedIslandLeaf, - right_leaf: &PlannedIslandLeaf, - ) -> Result { - let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty() - || right_leaf.ctx.tag_columns.is_empty()) - && !first_leaf - .ctx - .tag_columns - .iter() - .chain(&right_leaf.ctx.tag_columns) - .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL); - let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns( - left.schema(), - right_leaf.plan.schema(), - &first_leaf.ctx, - &right_leaf.ctx, - only_join_time_index, - &None, - )?; - - if let (Some(left_time_index_column), Some(right_time_index_column)) = ( - first_leaf.ctx.time_index_column.clone(), - right_leaf.ctx.time_index_column.clone(), - ) { - left_keys.insert(left_time_index_column); - right_keys.insert(right_time_index_column); - } - - LogicalPlanBuilder::from(left) - .join_detailed( - right_leaf.plan.clone(), - JoinType::Inner, - ( - left_keys - .into_iter() - .map(|name| Column::new(Some(first_leaf.alias.clone()), name)) - .collect::>(), - right_keys - .into_iter() - .map(|name| Column::new(Some(right_leaf.alias.clone()), name)) - .collect::>(), - ), - force_empty_join.then_some(lit(false)), - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - fn build_binary_island_field_exprs( - expr: &IslandExpr, - leaves: &[PlannedIslandLeaf], - schema: &DFSchemaRef, - ) -> Result { - match expr { - IslandExpr::VectorLeaf(id) => { - let leaf = &leaves[*id]; - let exprs = leaf - .ctx - .field_columns - .iter() - .map(|field| { - schema - .qualified_field_with_name(Some(&leaf.alias), field) - .context(DataFusionPlanningSnafu) - .map(|field| DfExpr::Column(field.into())) - }) - .collect::>>()?; - let names = leaf - .ctx - .field_columns - .iter() - .map(|field| format!("{}.{}", leaf.display_table, field)) - .collect(); - Ok(IslandFieldExprs { - exprs, - names, - scalar: false, - }) - } - IslandExpr::Scalar(expr) => Ok(IslandFieldExprs { - exprs: vec![expr.clone()], - names: vec![expr.schema_name().to_string()], - scalar: true, - }), - IslandExpr::Unary { input } => { - let input = Self::build_binary_island_field_exprs(input, leaves, schema)?; - let mut exprs = Vec::with_capacity(input.exprs.len()); - let mut names = Vec::with_capacity(input.names.len()); - for (expr, name) in input.exprs.into_iter().zip(input.names) { - exprs.push(DfExpr::Negative(Box::new(expr))); - names.push(format!("-{name}")); - } - Ok(IslandFieldExprs { - exprs, - names, - scalar: input.scalar, - }) - } - IslandExpr::Binary { op, lhs, rhs } => { - let same_leaf = match (&**lhs, &**rhs) { - (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right)) - if left == right => - { - Some(*left) - } - _ => None, - }; - let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?; - let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?; - let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let scalar = lhs.scalar && rhs.scalar; - let op = op.to_string(); - - let (exprs, names) = match (lhs.scalar, rhs.scalar) { - (true, true) => { - let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?; - let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]); - (vec![expr], vec![name]) - } - (true, false) => { - let mut exprs = Vec::with_capacity(rhs.exprs.len()); - let mut names = Vec::with_capacity(rhs.names.len()); - for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) { - exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?); - names.push(format!("{} {op} {rhs_name}", lhs.names[0])); - } - (exprs, names) - } - (false, true) => { - let mut exprs = Vec::with_capacity(lhs.exprs.len()); - let mut names = Vec::with_capacity(lhs.names.len()); - for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) { - exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?); - names.push(format!("{lhs_name} {op} {}", rhs.names[0])); - } - (exprs, names) - } - (false, false) => { - let mut exprs = Vec::new(); - let mut names = Vec::new(); - for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs - .exprs - .into_iter() - .zip(rhs.exprs) - .zip(lhs.names.into_iter().zip(rhs.names)) - .enumerate() - { - if let Some(leaf) = same_leaf { - let field = leaves[leaf] - .ctx - .field_columns - .get(idx) - .cloned() - .unwrap_or_else(|| lhs_name.clone()); - lhs_name = format!("lhs.{field}"); - rhs_name = format!("rhs.{field}"); - } - exprs.push(expr_builder(lhs_expr, rhs_expr)?); - names.push(format!("{lhs_name} {op} {rhs_name}")); - } - (exprs, names) - } - }; - - Ok(IslandFieldExprs { - exprs, - names, - scalar, - }) - } - } - } - - fn project_binary_island( - &mut self, - input: LogicalPlan, - base_alias: &TableReference, - base_ctx: &PromPlannerContext, - field_exprs: IslandFieldExprs, - ) -> Result { - self.ctx = base_ctx.clone(); - - let schema = input.schema(); - let non_field_exprs = base_ctx - .tag_columns - .iter() - .chain(base_ctx.time_index_column.iter()) - .map(|column| { - schema - .qualified_field_with_name(Some(base_alias), column) - .context(DataFusionPlanningSnafu) - .map(|field| DfExpr::Column(field.into())) - }); - let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid) - .into_iter() - .map(Ok); - - self.ctx.field_columns = field_exprs.names; - let field_exprs = field_exprs - .exprs - .into_iter() - .zip(self.ctx.field_columns.iter()) - .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::, name)))); - - let project_exprs = non_field_exprs - .chain(tsid_expr) - .chain(field_exprs) - .collect::>>()?; - - let plan = LogicalPlanBuilder::from(input) - .project(project_exprs) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - self.ctx.table_name = None; - self.ctx.schema_name = None; - - Ok(plan) - } - async fn prom_binary_expr_to_plan( &mut self, query_engine_state: &QueryEngineState, @@ -1401,7 +1231,7 @@ impl PromPlanner { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; let rhs_is_histogram = Self::field_column_is_native_histogram(&input_schema, col); - let rhs = DfExpr::Column(col.into()); + let rhs = DfExpr::Column(Column::from_name(col)); let mut binary_expr = match Self::native_histogram_binary_expr( *op, expr.clone(), @@ -1468,7 +1298,7 @@ impl PromPlanner { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; let lhs_is_histogram = Self::field_column_is_native_histogram(&input_schema, col); - let lhs = DfExpr::Column(col.into()); + let lhs = DfExpr::Column(Column::from_name(col)); let mut binary_expr = match Self::native_histogram_binary_expr( *op, lhs.clone(), @@ -1595,6 +1425,18 @@ impl PromPlanner { } } else if right_is_empty_metric && !left_is_empty_metric { self.ctx = left_context.clone(); + } else if rhs.value_type() == ValueType::Scalar + && lhs.value_type() == ValueType::Vector + { + // A computed scalar on the right (`time()`, `scalar(...)`) has no labels; + // the result keeps the labels of the vector on the left. + self.ctx = left_context.clone(); + } + // An operand without a table name, such as a comparison filter's projection, is + // joined under the `""` qualifier; resolve the projected columns against it. + // Otherwise a time index present on both sides is ambiguous. + if self.ctx.table_name.is_none() { + self.ctx.table_name = Some(String::new()); } // Computed scalars reach this join path instead of the literal projection paths. // Broadcast them for arithmetic in the same way as literal scalars. @@ -1634,6 +1476,16 @@ impl PromPlanner { // `vector()` uses EmptyMetric and keeps GreptimeDB's timestamp broadcast. let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric; + let only_join_time_index = lhs.value_type() == ValueType::Scalar + || rhs.value_type() == ValueType::Scalar + || has_empty_metric_operand + || ((left_context.tag_columns.is_empty() + || right_context.tag_columns.is_empty()) + && !left_context + .tag_columns + .iter() + .chain(&right_context.tag_columns) + .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)); let join_plan = self.join_on_non_field_columns( left_input, right_input, @@ -1641,21 +1493,34 @@ impl PromPlanner { right_table_ref.clone(), left_time_index_column, right_time_index_column, - lhs.value_type() == ValueType::Scalar - || rhs.value_type() == ValueType::Scalar - || has_empty_metric_operand - || ((left_context.tag_columns.is_empty() - || right_context.tag_columns.is_empty()) - && !left_context - .tag_columns - .iter() - .chain(&right_context.tag_columns) - .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)), + only_join_time_index, modifier, &left_context, &right_context, )?; let join_plan_schema = join_plan.schema().clone(); + // The matching modifier derives the result labels instead of the operands keeping + // their own tag set. An operand that is broadcast rather than matched (a scalar, + // or a vector without tags) keeps the other side's labels as before. + let result_labels = if only_join_time_index { + None + } else { + Self::binary_result_labels(&left_context, &right_context, modifier) + .map(|labels| { + Self::binary_result_label_projection( + &join_plan_schema, + &left_table_ref, + &right_table_ref, + &left_context, + &right_context, + labels, + ) + }) + .transpose()? + }; + if let Some(labels) = &result_labels { + labels.apply(&mut self.ctx); + } let promql_annotations = self.promql_annotations.clone(); // These predicates always pass; they only evaluate otherwise-discarded pairs // while collecting annotations. @@ -1776,10 +1641,18 @@ impl PromPlanner { ), }; project_context.field_columns = project_field_columns; - self.project_binary_join_side(filtered, project_table_ref, &project_context) + self.project_binary_join_side( + filtered, + project_table_ref, + &project_context, + result_labels.as_ref(), + ) } else { - let projected = - self.projection_for_each_field_column(join_plan, bin_expr_builder)?; + let projected = self.projection_for_each_field_column_with_labels( + join_plan, + result_labels.as_ref(), + bin_expr_builder, + )?; let preserve_any_value = Self::field_columns_are_alternative_samples( projected.schema(), &self.ctx.field_columns, @@ -1867,6 +1740,7 @@ impl PromPlanner { input: LogicalPlan, table_ref: &TableReference, context: &PromPlannerContext, + result_labels: Option<&BinaryResultLabels>, ) -> Result { let schema = input.schema(); @@ -1891,20 +1765,27 @@ impl PromPlanner { project_exprs.push(DfExpr::Column(field_col)); } - // Project tag columns from the chosen side. - for tag_column in &context.tag_columns { - let tag_col = schema - .qualified_field_with_name(Some(table_ref), tag_column) - .context(DataFusionPlanningSnafu)? - .into(); - project_exprs.push(DfExpr::Column(tag_col)); + // Project tag columns: the labels the matching modifier derived, or the chosen side's. + match result_labels { + Some(labels) => project_exprs.extend(labels.exprs.iter().cloned()), + None => { + for tag_column in &context.tag_columns { + let tag_col = schema + .qualified_field_with_name(Some(table_ref), tag_column) + .context(DataFusionPlanningSnafu)? + .into(); + project_exprs.push(DfExpr::Column(tag_col)); + } + } } // Preserve `__tsid` if present, so it can still be used internally downstream. It's // stripped from the final output anyway. - if let Some(tsid_col) = - Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid) - { + let tsid_col = match result_labels { + Some(labels) => labels.tsid_projection(Some(table_ref.clone())), + None => Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid), + }; + if let Some(tsid_col) = tsid_col { project_exprs.push(tsid_col); } @@ -1919,6 +1800,9 @@ impl PromPlanner { self.ctx = context.clone(); self.ctx.table_name = None; self.ctx.schema_name = None; + if let Some(labels) = result_labels { + labels.apply(&mut self.ctx); + } Ok(plan) } @@ -1969,6 +1853,85 @@ impl PromPlanner { Ok(plan) } + /// The offset of a selector in milliseconds. A positive offset selects samples from an earlier + /// time and moves them forward into the evaluation timeline. + fn offset_millis(offset: &Option) -> Millisecond { + match offset { + Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, + Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), + None => 0, + } + } + + /// Returns the first absolute step multiple strictly after `start - offset - range` + /// through the last parent evaluation instant minus `offset`. Requires a positive step. + /// Returns `None` for an empty window; an unrepresentable nonempty bound is an error. + fn subquery_child_window( + start: Millisecond, + end: Millisecond, + interval: Millisecond, + offset: Millisecond, + range: Millisecond, + step: Millisecond, + ) -> Result> { + let window_start = i128::from(start) - i128::from(offset) - i128::from(range); + let grid_start = (window_start.div_euclid(i128::from(step)) + 1) * i128::from(step); + let grid_end = if interval > 0 { + let steps = (i128::from(end) - i128::from(start)) / i128::from(interval); + i128::from(start) + steps * i128::from(interval) + } else { + i128::from(end) + } - i128::from(offset); + // Check emptiness before converting bounds to i64. + if grid_start > grid_end { + return Ok(None); + } + let bound = |name: &str, value: i128| { + i64::try_from(value).map_err(|_| { + SubqueryTimestampOutOfRangeSnafu { + timestamp: format!( + "subquery child window {name} = {value} ms (start = {start} ms, \ + end = {end} ms, interval = {interval} ms, offset = {offset} ms, \ + range = {range} ms, step = {step} ms)" + ), + } + .build() + }) + }; + Ok(Some(( + bound("grid start", grid_start)?, + bound("grid end", grid_end)?, + ))) + } + + /// The columns that identify one series, which is the series key expected by the PromQL plan + /// nodes that hold exactly one series per input batch. + fn series_key_columns(&self) -> Vec { + if self.ctx.use_tsid { + vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] + } else { + self.ctx.tag_columns.clone() + } + } + + /// Keep replay and series division on the same effective keys when a call rewrites labels. + fn series_key_columns_for_schema(&self, schema: &DFSchemaRef) -> Vec { + let has_tsid = schema.fields().iter().any(|field| { + field.name() == DATA_SCHEMA_TSID_COLUMN_NAME + && field.data_type() == &ArrowDataType::UInt64 + }); + if has_tsid { + vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] + } else { + self.ctx + .tag_columns + .iter() + .filter(|name| schema.has_column_with_unqualified_name(name)) + .cloned() + .collect() + } + } + async fn prom_vector_selector_to_plan( &mut self, vector_selector: &VectorSelector, @@ -1978,20 +1941,37 @@ impl PromPlanner { name, offset, matchers, - at: _, + at, } = vector_selector; let matchers = self.preprocess_label_matchers(matchers, name)?; 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 offset_ms = Self::offset_millis(offset); + // `@` anchors the sample window at a fixed timestamp: the selector selects its samples + // around the anchor once, instead of following the outer evaluation grid. See + // [`Self::at_modifier_offset`]. + let at_offset = self.at_modifier_offset(at, offset)?; + let grid_start = self.ctx.start; + let grid_end = self.ctx.end; + let normalize = match at_offset { + Some(at_offset) => { + // Select the anchored samples at the start of the evaluation, with the offset that + // re-anchors the selector. + // The planner is single-use (one `EvalStmt` produces one plan), so an error + // below aborts the whole planning and `ctx.end` needs no restore-on-error. + self.ctx.end = grid_start; + let plan = self + .selector_to_series_normalize_plan(at_offset, matchers, false) + .await?; + self.ctx.end = grid_end; + plan + } + None => { + self.selector_to_series_normalize_plan(offset_ms, matchers, false) + .await? + } }; - let normalize = self - .selector_to_series_normalize_plan(offset, matchers, false) - .await?; let time_index_column = self.ctx .time_index_column @@ -2020,47 +2000,9 @@ impl PromPlanner { DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) }) .collect::>(); - // `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)), - }); + // The time index still holds the raw sample timestamp here, which is what + // `timestamp()` reports regardless of `offset` and `@`. + let sample_time = Self::timestamp_seconds_expr(&time_index_column, normalize.schema())?; project_exprs.push(sample_time.alias(×tamp_value_column)); let normalize = LogicalPlanBuilder::from(normalize) .project(project_exprs) @@ -2073,59 +2015,92 @@ impl PromPlanner { }; let field_column = self.ctx.field_columns.first().cloned(); - let manipulate = InstantManipulate::new( - self.ctx.start, - 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()] - } else { - self.ctx.tag_columns.clone() - }, - field_column, - normalize, - ); - let manipulate = LogicalPlan::Extension(Extension { - node: Arc::new(manipulate), - }); + let series_key_columns = self.series_key_columns(); + let manipulate = match at_offset { + Some(at_offset) => { + // Select the anchored sample once, then report it at every step of the outer + // grid. The samples keep their native anchor-time timestamps, so the manipulate + // must shift them onto the evaluation timeline with the rewritten offset. + let anchored = InstantManipulate::new( + grid_start, + grid_start, + self.ctx.lookback_delta, + self.ctx.interval, + at_offset, + time_index_column.clone(), + series_key_columns, + field_column, + normalize, + ); + self.replay_over_grid( + LogicalPlan::Extension(Extension { + node: Arc::new(anchored), + }), + grid_start, + grid_end, + time_index_column, + ) + } + None => LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + grid_start, + grid_end, + self.ctx.lookback_delta, + self.ctx.interval, + offset_ms, + time_index_column, + series_key_columns, + field_column, + normalize, + )), + }), + }; if let Some(timestamp_value_column) = timestamp_value_column { - self.create_timestamp_func_plan(manipulate, ×tamp_value_column) + self.create_timestamp_func_plan(manipulate, ident(timestamp_value_column)) } else { Ok(manipulate) } } - /// Builds a projection plan for the PromQL `timestamp()` function. - /// Projects the time index column as the value column for each row. + /// Converts the timestamp column `column` into PromQL seconds, truncated to milliseconds. + fn timestamp_seconds_expr(column: &str, schema: &DFSchema) -> Result { + let column = DfExpr::Column(Column::from_name(column)); + let ArrowDataType::Timestamp(_, timezone) = + column.get_type(schema).context(DataFusionPlanningSnafu)? + else { + unreachable!("time index is a timestamp") + }; + let millis = column + .cast_to( + &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, timezone), + schema, + ) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Int64, schema) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Float64, schema) + .context(DataFusionPlanningSnafu)?; + Ok(DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(millis), + op: Operator::Divide, + right: Box::new(lit(1000.0)), + })) + } + + /// Builds a projection plan for the PromQL `timestamp()` function, which reports + /// `timestamp_value` as the value of each row, along with the original tag and time index + /// columns. /// - /// # Arguments - /// * `input` - Input [`LogicalPlan`] after instant-vector selection. - /// * `timestamp_value_column` - Private column containing each selected sample's timestamp. - /// - /// # Returns - /// Returns a [`Result`] where the resulting logical plan projects the timestamp - /// column as the value column, along with the original tag and time index columns. - /// - /// # Timestamp vs. Time Function - /// - /// - **Timestamp Function (`timestamp()`)**: In PromQL, the `timestamp()` function returns the - /// timestamp (time index) of each sample as the value column. - /// - /// - **Time Function (`time()`)**: The `time()` function returns the evaluation time of the query - /// as a scalar value. - /// - /// # Side Effects /// Updates the planner context's field columns to the timestamp column name. - /// fn create_timestamp_func_plan( &mut self, input: LogicalPlan, - timestamp_value_column: &str, + timestamp_value: DfExpr, ) -> Result { - let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN); + // A row whose fields are all NULL holds no sample, so it must not get a timestamp. The + // check reads the input fields, which the projection below replaces. + let has_sample = self.create_empty_values_filter_expr(true)?; + let time_expr = timestamp_value.alias(DEFAULT_FIELD_COLUMN); self.ctx.field_columns = vec![time_expr.schema_name().to_string()]; let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2); project_exprs.push(self.create_time_index_column_expr()?); @@ -2133,6 +2108,8 @@ impl PromPlanner { project_exprs.extend(self.create_tag_column_exprs()?); LogicalPlanBuilder::from(input) + .filter(has_sample) + .context(DataFusionPlanningSnafu)? .project(project_exprs) .context(DataFusionPlanningSnafu)? .build() @@ -2148,46 +2125,108 @@ impl PromPlanner { name, offset, matchers, - .. + at, } = vs; let matchers = self.preprocess_label_matchers(matchers, name)?; 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, - }; + let offset_ms = Self::offset_millis(offset); + + // `@` anchors the range selector's window at a fixed timestamp, so the same window is fed + // to the enclosing function at every evaluation step. See [`Self::at_modifier_offset`]. + let at_offset = self.at_modifier_offset(at, offset)?; + let grid_start = self.ctx.start; + let grid_end = self.ctx.end; // Some functions like rate may require special fields in the RangeManipulate plan // so we can't skip RangeManipulate. - let normalize = match self.setup_context().await? { - Some(empty_plan) => empty_plan, + let (normalize, at_offset) = match self.setup_context().await? { + // An empty metric does not contain any sample, so anchoring cannot change the result. + // The manipulate below folds the empty input with `offset_ms`, and the recorded fold + // offset agrees with it instead of with the anchor the window cannot use. + Some(empty_plan) => { + self.ctx.range_fold_offset = Some(offset_ms); + (empty_plan, None) + } None => { - self.selector_to_series_normalize_plan(offset, matchers, true) - .await? + let normalize = match at_offset { + Some(at_offset) => { + // Fold the anchored window once, at the start of the evaluation. + // Single-use planner: an error below aborts planning, so `ctx.end` + // needs no restore-on-error. + self.ctx.end = grid_start; + let plan = self + .selector_to_series_normalize_plan(at_offset, matchers, true) + .await?; + self.ctx.end = grid_end; + plan + } + None => { + self.selector_to_series_normalize_plan(offset_ms, matchers, true) + .await? + } + }; + // Samples are shifted onto the evaluation timeline with the very same offset while + // the window is folded. Record it so that the function above the selector can + // recover the evaluation instant of the folded window; see + // [`Self::create_range_eval_ts_expr`]. + self.ctx.range_fold_offset = Some(at_offset.unwrap_or(offset_ms)); + (normalize, at_offset) } }; - let manipulate = RangeManipulate::new( - 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 - .time_index_column - .clone() - .expect("time index should be set in `setup_context`"), - self.ctx.field_columns.clone(), - normalize, - ) - .context(DataFusionPlanningSnafu)?; + let time_index_column = self + .ctx + .time_index_column + .clone() + .expect("time index should be set in `setup_context`"); + let manipulate = match at_offset { + Some(at_offset) => { + // Fold the anchored window once, then report it at every step of the outer + // grid. The samples keep their native anchor-time timestamps, so the manipulate + // must shift them onto the evaluation timeline with the rewritten offset. + let anchored = RangeManipulate::new( + grid_start, + grid_start, + self.ctx.interval, + at_offset, + // TODO(ruihang): convert via Timestamp datatypes to support different time units + range_ms, + time_index_column.clone(), + self.ctx.field_columns.clone(), + normalize, + ) + .context(DataFusionPlanningSnafu)?; + self.replay_over_grid( + LogicalPlan::Extension(Extension { + node: Arc::new(anchored), + }), + grid_start, + grid_end, + time_index_column, + ) + } + None => { + let manipulate = RangeManipulate::new( + grid_start, + grid_end, + self.ctx.interval, + offset_ms, + // TODO(ruihang): convert via Timestamp datatypes to support different time units + range_ms, + time_index_column, + self.ctx.field_columns.clone(), + normalize, + ) + .context(DataFusionPlanningSnafu)?; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(manipulate), - })) + LogicalPlan::Extension(Extension { + node: Arc::new(manipulate), + }) + } + }; + + Ok(manipulate) } async fn prom_call_expr_to_plan( @@ -2213,9 +2252,31 @@ impl PromPlanner { // transform function arguments let args = self.create_function_args(&args.args)?; + // Only a vector selector keeps the timestamps of its samples. Any other expression + // produces samples at the evaluation time, which is what `timestamp()` reports for it. + let mut timestamp_arg = args.input.as_ref(); + while let Some(PromExpr::Paren(ParenExpr { expr })) = timestamp_arg { + timestamp_arg = Some(expr); + } + let timestamp_of_selector = + func.name == "timestamp" && matches!(timestamp_arg, Some(PromExpr::VectorSelector(_))); let input = if let Some(prom_expr) = &args.input { - self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state) - .await? + let input = self + .prom_expr_to_plan_inner(prom_expr, timestamp_of_selector, query_engine_state) + .await?; + if func.name == "timestamp" && !timestamp_of_selector { + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let eval_time = Self::timestamp_seconds_expr(&time_index_column, input.schema())?; + self.create_timestamp_func_plan(input, eval_time)? + } else { + input + } } else { self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); self.ctx.reset_table_name_and_schema(); @@ -2236,11 +2297,18 @@ impl PromPlanner { ), }) }; + // The input plan records the fold offset of the range selector it is built from. Take it + // here, so that the offset of one input cannot leak into another call, and pass it to + // `create_function_expr`: the function that reads it (`predict_linear`) then depends on an + // argument instead of on planner state written by the selector below it. + let range_fold_offset = self.ctx.range_fold_offset.take(); + let input_tag_columns = self.ctx.tag_columns.clone(); let (mut func_exprs, new_tags) = self.create_function_expr( func, args.literals.clone(), input.schema(), query_engine_state, + range_fold_offset, )?; func_exprs.insert(0, self.create_time_index_column_expr()?); func_exprs.extend_from_slice(&self.create_tag_column_exprs()?); @@ -2286,13 +2354,31 @@ impl PromPlanner { _ => builder, }; + // Rewriting a label the input series already have can map several of them onto the same + // label set, which PromQL rejects. A new label keeps the series distinct. + let may_duplicate_label_sets = matches!(func.name, "label_join" | "label_replace") + && new_tags.iter().any(|tag| input_tag_columns.contains(tag)); + // Update context tags after building plan // We can't push them before planning, because they won't exist until projection. for tag in new_tags { self.ctx.tag_columns.push(tag); } - let plan = builder.build().context(DataFusionPlanningSnafu)?; + let mut plan = builder.build().context(DataFusionPlanningSnafu)?; + if may_duplicate_label_sets { + let labels = self.ctx.tag_columns.clone(); + plan = Self::assert_unique_match_group( + plan, + labels + .iter() + .map(|label| DfExpr::Column(Column::from_name(label))) + .collect(), + labels, + self.create_time_index_column_expr()?, + MatchGroupViolation::DuplicateLabelSet, + )?; + } common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}"); Ok(plan) @@ -2307,7 +2393,7 @@ impl PromPlanner { let expr = &ext_expr.expr; let children = expr.children(); let plan = self - .prom_expr_to_plan(&children[0], query_engine_state) + .prom_root_expr_to_plan(&children[0], query_engine_state) .await?; // Wrapper for the explanation/analyze of the existing plan // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain @@ -2424,7 +2510,7 @@ impl PromPlanner { async fn selector_to_series_normalize_plan( &mut self, - offset: &Option, + offset_duration: Millisecond, label_matchers: Matchers, is_range_selector: bool, ) -> Result { @@ -2434,11 +2520,6 @@ impl PromPlanner { let table_schema = table_scan.schema(); // make filter exprs - let offset_duration = match offset { - Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, - Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), - 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, table_schema)? @@ -2537,11 +2618,7 @@ impl PromPlanner { } // make sort plan - let series_key_columns = if self.ctx.use_tsid { - vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] - } else { - self.ctx.tag_columns.clone() - }; + let series_key_columns = self.series_key_columns(); let sort_exprs = if self.ctx.use_tsid { vec![ @@ -2696,7 +2773,7 @@ impl PromPlanner { // collect remaining fields and convert to col expr let mut exprs = all_fields .into_iter() - .map(|c| DfExpr::Column(Column::from(c))) + .map(|c| DfExpr::Column(Column::from_name(c))) .collect::>(); // add timestamp column @@ -2728,28 +2805,21 @@ impl PromPlanner { // for it, or the column exists but is NULL on that row — the latter // is the norm for logical metrics sharing a physical table, which // holds the union of their label columns. + let mut null_label = None; let col = if let Some(column_name) = column_name { let column = DfExpr::Column(Column::from_name(&column_name)); let field = table_schema .index_of_column_by_name(None, &column_name) .map(|index| table_schema.field(index)); if accepts_empty - && let Some(data_type) = field - .filter(|field| { - field.is_nullable() - && Self::string_value_data_type(field.data_type()).is_some() - }) - .map(|field| field.data_type()) - { - let empty = Self::string_scalar_value(data_type, Some(String::new())) - .expect("nullable label has a string type"); - DfExpr::ScalarFunction(ScalarFunction { - func: coalesce(), - args: vec![column, DfExpr::Literal(empty, None)], + && field.is_some_and(|field| { + field.is_nullable() + && Self::string_value_data_type(field.data_type()).is_some() }) - } else { - column + { + null_label = Some(column.clone().is_null()); } + column } else { DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None) }; @@ -2802,6 +2872,10 @@ impl PromPlanner { } } }; + let expr = match null_label { + Some(null_label) => null_label.or(expr), + None => expr, + }; exprs.push(expr); } @@ -3214,130 +3288,6 @@ impl PromPlanner { Ok(out) } - fn ensure_tag_columns_available( - &self, - plan: LogicalPlan, - required_tags: &BTreeSet, - ) -> Result { - if required_tags.is_empty() { - return Ok(plan); - } - - struct Rewriter { - required_tags: BTreeSet, - } - - impl TreeNodeRewriter for Rewriter { - type Node = LogicalPlan; - - fn f_up( - &mut self, - node: Self::Node, - ) -> datafusion_common::Result> { - match node { - LogicalPlan::TableScan(scan) => { - let schema = scan.source.schema(); - let mut projection = match scan.projection.clone() { - Some(p) => p, - None => { - // Scanning all columns already covers required tags. - return Ok(Transformed::no(LogicalPlan::TableScan(scan))); - } - }; - - let mut changed = false; - for tag in &self.required_tags { - if let Some((idx, _)) = schema - .fields() - .iter() - .enumerate() - .find(|(_, field)| field.name() == tag) - && !projection.contains(&idx) - { - projection.push(idx); - changed = true; - } - } - - if !changed { - return Ok(Transformed::no(LogicalPlan::TableScan(scan))); - } - - projection.sort_unstable(); - projection.dedup(); - - let new_scan = - TableScanBuilder::new(scan.table_name.clone(), scan.source.clone()) - .with_projection(Some(projection)) - .with_filters(scan.filters) - .with_fetch(scan.fetch) - .build()?; - Ok(Transformed::yes(LogicalPlan::TableScan(new_scan))) - } - LogicalPlan::Projection(proj) => { - let input_schema = proj.input.schema(); - - let existing = proj - .schema - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect::>(); - - let mut expr = proj.expr.clone(); - let mut has_changed = false; - for tag in &self.required_tags { - if existing.contains(tag.as_str()) { - continue; - } - - if let Some(idx) = input_schema.index_of_column_by_name(None, tag) { - expr.push(DfExpr::Column(Column::from( - input_schema.qualified_field(idx), - ))); - has_changed = true; - } - } - - if !has_changed { - return Ok(Transformed::no(LogicalPlan::Projection(proj))); - } - - let new_proj = Projection::try_new(expr, proj.input)?; - Ok(Transformed::yes(LogicalPlan::Projection(new_proj))) - } - other => Ok(Transformed::no(other)), - } - } - } - - let mut rewriter = Rewriter { - required_tags: required_tags.clone(), - }; - let rewritten = plan - .rewrite(&mut rewriter) - .context(DataFusionPlanningSnafu)?; - Ok(rewritten.data) - } - - fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) { - let time_index = self.ctx.time_index_column.as_deref(); - let field_columns = self.ctx.field_columns.iter().collect::>(); - - let mut tags = schema - .fields() - .iter() - .map(|f| f.name()) - .filter(|name| Some(name.as_str()) != time_index) - .filter(|name| !field_columns.contains(name)) - .filter(|name| !is_metric_engine_internal_column(name)) - .cloned() - .collect::>(); - tags.sort_unstable(); - tags.dedup(); - self.ctx.tag_columns = tags; - } - /// Setup [PromPlannerContext]'s state fields. /// /// Returns a logical plan for an empty metric. @@ -3462,6 +3412,7 @@ impl PromPlanner { float_field: &str, histogram_field: &str, input_schema: &DFSchemaRef, + range_fold_offset: Option, ) -> Result>> { let returns_histogram = matches!( func.name, @@ -3501,6 +3452,14 @@ impl PromPlanner { Box::new(other_input_exprs[0].clone()), ArrowDataType::Int64, )); + // Same evaluation instant as in the non-mixed path; it follows the other inputs so the + // float UDF below is called as `predict_linear(ts_range, value_range, t, eval_ts)`. + // The regression is only defined for a folded window, so a missing fold offset is an + // error instead of a default. + other_input_exprs.push_back(self.create_range_eval_ts_expr( + range_fold_offset.context(ExpectRangeSelectorSnafu)?, + input_schema, + )?); } let timestamp_range = DfExpr::Column(Column::from_name( @@ -3559,7 +3518,13 @@ impl PromPlanner { .alias(histogram_field), ] } else { - let display_name = float_expr.schema_name().to_string(); + let display_name = if func.name == "predict_linear" { + // The evaluation instant is the private last argument of the mixed float UDF + // call; keep it out of the output column name like on the non-mixed path. + Self::name_without_last_arg(&float_expr) + } else { + float_expr.schema_name().to_string() + }; self.ctx.field_columns = vec![display_name.clone()]; vec![float_expr.alias(display_name)] }; @@ -3577,6 +3542,7 @@ impl PromPlanner { other_input_exprs: Vec, input_schema: &DFSchemaRef, query_engine_state: &QueryEngineState, + range_fold_offset: Option, ) -> Result<(Vec, Vec)> { // TODO(ruihang): check function args list let mut other_input_exprs: VecDeque = other_input_exprs.into(); @@ -3589,6 +3555,7 @@ impl PromPlanner { &float_field, &histogram_field, input_schema, + range_fold_offset, )? { return Ok((exprs, vec![])); @@ -3790,6 +3757,16 @@ impl PromPlanner { Box::new(other_input_exprs[0].clone()), ArrowDataType::Int64, )); + // The prediction starts at the evaluation instant of the step, which the + // window's fold offset recovers from the row's time index; see + // [`Self::create_range_eval_ts_expr`]. It is appended last, so the UDF is + // called as `predict_linear(ts_range, value_range, t, eval_ts)`. The + // regression is only defined for a folded window, so a missing fold offset + // is an error instead of a default. + other_input_exprs.push_back(self.create_range_eval_ts_expr( + range_fold_offset.context(ExpectRangeSelectorSnafu)?, + input_schema, + )?); ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf())) } } @@ -3904,6 +3881,7 @@ impl PromPlanner { let (concat_expr, dst_label) = Self::build_concat_labels_expr( &mut other_input_exprs, &self.ctx, + input_schema, query_engine_state, )?; @@ -3925,8 +3903,8 @@ impl PromPlanner { } "label_replace" => { self.ctx.use_tsid = false; - if let Some((replace_expr, dst_label)) = self - .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)? + if let Some((replace_expr, dst_label)) = + self.build_regexp_replace_label_expr(&mut other_input_exprs, input_schema)? { // Reserve the current field columns except the `dst_label`. for value in &self.ctx.field_columns { @@ -3936,10 +3914,8 @@ impl PromPlanner { } } - ensure!( - !self.ctx.tag_columns.contains(&dst_label), - SameLabelSetSnafu - ); + // Remove it from tag columns if exists to avoid duplicated column names + self.ctx.tag_columns.retain(|tag| *tag != dst_label); new_tags.push(dst_label); // Add the new label expr to evaluate exprs.push(replace_expr); @@ -4162,7 +4138,17 @@ impl PromPlanner { exprs = exprs .into_iter() .map(|expr| { - let display_name = expr.schema_name().to_string(); + // `predict_linear` appends its private evaluation instant as the last argument + // of the UDF call; the output column is named after the call without it, so + // the injected expression stays out of the user-visible schema. The + // native-histogram drop UDF takes no such argument. + let display_name = if func.name == "predict_linear" + && !all_field_columns_are_native_histogram_ranges + { + Self::name_without_last_arg(&expr) + } else { + expr.schema_name().to_string() + }; new_field_columns.push(display_name.clone()); Ok(expr.alias(display_name)) }) @@ -4238,7 +4224,7 @@ impl PromPlanner { fn build_regexp_replace_label_expr( &self, other_input_exprs: &mut VecDeque, - query_engine_state: &QueryEngineState, + input_schema: &DFSchemaRef, ) -> Result> { // label_replace(vector, dst_label, replacement, src_label, regex) let dst_label = match other_input_exprs.pop_front() { @@ -4274,72 +4260,64 @@ impl PromPlanner { .fail()?, }; - // Validate the regex before using it + // Like Prometheus, match the whole source value. A series whose source value matches gets + // `dst_label` set to the expanded replacement; any other series is left unchanged. // doc: https://prometheus.io/docs/prometheus/latest/querying/functions/#label_replace - regex::Regex::new(®ex).map_err(|_| { - InvalidRegularExpressionSnafu { - regex: regex.clone(), - } - .build() - })?; + let anchored = format!("^(?s:{regex})$"); + let compiled = regex::Regex::new(&anchored) + .map_err(|_| InvalidRegularExpressionSnafu { regex }.build())?; + let dst_exists = self.ctx.tag_columns.contains(&dst_label); - // If the src_label exists and regex is empty, keep everything unchanged. - if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() { - return Ok(None); - } - - // If the src_label doesn't exists, and + // A missing source label reads as the empty string for every series, so the result is + // the same for all of them and is decided here. if !self.ctx.tag_columns.contains(&src_label) { - if replacement.is_empty() { - // the replacement is empty, keep everything unchanged. + let Some(captures) = compiled.captures("") else { return Ok(None); - } else { - // the replacement is not empty, always adds dst_label with replacement value. + }; + let mut value = String::new(); + captures.expand(&replacement, &mut value); + if value.is_empty() { + // Setting a label to the empty string removes it, which is a no-op for a new + // label. + if !dst_exists { + return Ok(None); + } return Ok(Some(( - // alias literal `replacement` as dst_label - lit(replacement).alias(&dst_label), + lit(ScalarValue::Utf8(None)).alias(&dst_label), dst_label, ))); } + return Ok(Some((lit(value).alias(&dst_label), dst_label))); } - // Preprocess the regex: - // https://github.com/prometheus/prometheus/blob/d902abc50d6652ba8fe9a81ff8e5cce936114eba/promql/functions.go#L1575C32-L1575C37 - let regex = format!("^(?s:{regex})$"); + let src = Self::label_value_expr(&src_label, input_schema)?; + let matched = DfExpr::ScalarFunction(ScalarFunction { + func: datafusion_functions::regex::regexp_like(), + args: vec![src.clone(), lit(anchored.clone())], + }); + let replaced = Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction { + func: datafusion_functions::regex::regexp_replace(), + args: vec![src, lit(anchored), lit(replacement)], + })); + let unchanged = if dst_exists { + DfExpr::Column(Column::from_name(&dst_label)) + .cast_to(&ArrowDataType::Utf8, input_schema) + .context(DataFusionPlanningSnafu)? + } else { + lit(ScalarValue::Utf8(None)) + }; + let replace_expr = when(matched, replaced) + .otherwise(unchanged) + .context(DataFusionPlanningSnafu)?; - let session_state = query_engine_state.session_state(); - let func = session_state - .scalar_functions() - .get("regexp_replace") - .context(UnsupportedExprSnafu { - name: "regexp_replace", - })?; - - // regexp_replace(src_label, regex, replacement) - let args = vec![ - if src_label.is_empty() { - DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None) - } else { - DfExpr::Column(Column::from_name(src_label)) - }, - DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None), - DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None), - ]; - - Ok(Some(( - DfExpr::ScalarFunction(ScalarFunction { - func: func.clone(), - args, - }) - .alias(&dst_label), - dst_label, - ))) + Ok(Some((replace_expr.alias(&dst_label), dst_label))) } /// Build expr for `label_join` function fn build_concat_labels_expr( other_input_exprs: &mut VecDeque, ctx: &PromPlannerContext, + input_schema: &DFSchemaRef, query_engine_state: &QueryEngineState, ) -> Result<(DfExpr, String)> { // label_join(vector, dst_label, separator, src_label_1, src_label_2, ...) @@ -4371,17 +4349,19 @@ impl PromPlanner { let src_labels = other_input_exprs .iter() .map(|expr| { - // Cast source label into column or null literal + // `concat_ws` skips NULL arguments together with their separator, while an + // absent label joins as the empty string. match expr { DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => { if label.is_empty() { - Ok(DfExpr::Literal(ScalarValue::Null, None)) + FunctionInvalidArgumentSnafu { + fn_name: "label_join", + } + .fail() } else if available_columns.contains(label.as_str()) { - // Label exists in the table schema - Ok(DfExpr::Column(Column::from_name(label))) + Self::label_value_expr(label, input_schema) } else { - // Label doesn't exist, treat as empty string (null) - Ok(DfExpr::Literal(ScalarValue::Null, None)) + Ok(lit("")) } } other => UnexpectedPlanExprSnafu { @@ -4394,12 +4374,10 @@ impl PromPlanner { } }) .collect::>>()?; - ensure!( - !src_labels.is_empty(), - FunctionInvalidArgumentSnafu { - fn_name: "label_join" - } - ); + // Joining no labels yields the empty string, i.e. removes `dst_label`. + if src_labels.is_empty() { + return Ok((lit(ScalarValue::Utf8(None)).alias(&dst_label), dst_label)); + } let session_state = query_engine_state.session_state(); let func = session_state @@ -4413,15 +4391,37 @@ impl PromPlanner { args.extend(src_labels); Ok(( - DfExpr::ScalarFunction(ScalarFunction { + Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction { func: func.clone(), args, - }) + })) .alias(&dst_label), dst_label, )) } + /// The value of `label` as a string, where NULL (the series has no such label) reads as the + /// empty string, as in PromQL. + fn label_value_expr(label: &str, input_schema: &DFSchemaRef) -> Result { + let value = DfExpr::Column(Column::from_name(label)) + .cast_to(&ArrowDataType::Utf8, input_schema) + .context(DataFusionPlanningSnafu)?; + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![value, lit("")], + })) + } + + /// An empty label value means the label is absent in PromQL. Label functions represent it + /// as NULL, like a series that never had the label, so that both compare equal when label + /// sets are matched and neither is reported as a label. + fn empty_label_to_null(value: DfExpr) -> DfExpr { + DfExpr::ScalarFunction(ScalarFunction { + func: datafusion_functions::core::nullif(), + args: vec![value, lit("")], + }) + } + fn create_time_index_column_expr(&self) -> Result { Ok(DfExpr::Column(Column::from_name( self.ctx @@ -4431,6 +4431,66 @@ impl PromPlanner { ))) } + /// Builds the evaluation instant the window of the last planned range selector is folded for, + /// as a `Timestamp(Millisecond)` expression. + /// + /// The timestamp payload of a folded window is shifted onto the evaluation timeline by the + /// offset the window is folded with (`fold_offset`), while the time index column of a folded + /// row keeps the evaluation timestamp of its step. Adding the offset back yields the + /// evaluation instant on the payload timeline, which is where the regression of + /// `predict_linear` is centered: neither a plain window (which may end before the step, and + /// is additionally shifted by `offset` on the payload timeline) nor an `@`-anchored one + /// (whose end is the anchor, while the payload is shifted by `at_offset`) ends at the step it + /// is evaluated at. + /// + /// The sum is computed on the millisecond representation and cast back, so that the result + /// keeps the `Timestamp(Millisecond)` type the range functions declare for it. + fn create_range_eval_ts_expr( + &self, + fold_offset: Millisecond, + input_schema: &DFSchemaRef, + ) -> Result { + let eval_ts = self + .create_time_index_column_expr()? + .cast_to( + &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + input_schema, + ) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Int64, input_schema) + .context(DataFusionPlanningSnafu)?; + DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(eval_ts), + op: Operator::Plus, + right: Box::new(lit(fold_offset)), + }) + .cast_to( + &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + input_schema, + ) + .context(DataFusionPlanningSnafu) + } + + /// The name of `expr` without its last argument. + /// + /// A `predict_linear` call appends the evaluation instant of its window as a private last + /// argument ([`Self::create_range_eval_ts_expr`]). Naming the output column after the call + /// the user wrote, without that argument, keeps the injected expression out of the + /// user-visible schema; the expression itself keeps every argument it needs. + fn name_without_last_arg(expr: &DfExpr) -> String { + if let DfExpr::ScalarFunction(ScalarFunction { func, args }) = expr + && let Some((_, visible_args)) = args.split_last() + { + let visible = ScalarFunction { + func: func.clone(), + args: visible_args.to_vec(), + }; + return DfExpr::ScalarFunction(visible).schema_name().to_string(); + } + + expr.schema_name().to_string() + } + fn create_tag_column_exprs(&self) -> Result> { let mut result = Vec::with_capacity(self.ctx.tag_columns.len()); for tag in &self.ctx.tag_columns { @@ -4526,9 +4586,10 @@ impl PromPlanner { /// /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where: /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields - /// - `previous_field_expressions`: Original field expressions before aggregation. This is non-empty - /// only when the operation is `count_values`, as this operation requires preserving the original - /// values for grouping. + /// - `previous_field_expressions`: Field expressions naming the pre-aggregation values. This is + /// non-empty only when the operation is `count_values`, which groups by the sample value and + /// projects it as the generated label, so these expressions are passed through the same + /// formatting as that label (`prom_float_to_string`). /// fn create_aggregate_exprs( &mut self, @@ -4572,14 +4633,13 @@ impl PromPlanner { .collect::>>()?; // if the aggregator is `count_values`, it must be grouped by current fields. + // + // The grouping key is the *formatted* sample value, i.e. the same expression that + // produces the generated label below: PromQL groups by the value, and the label is + // that value in Prometheus' textual form (`strconv.FormatFloat(value, 'f', -1, 64)`), + // so grouping by the raw value would split samples that render to one label into + // several groups, each emitting the same label set for one timestamp. let prev_field_exprs = if op.id() == token::T_COUNT_VALUES { - let prev_field_exprs: Vec<_> = self - .ctx - .field_columns - .iter() - .map(|col| DfExpr::Column(Column::from_name(col))) - .collect(); - ensure!( self.ctx.field_columns.len() == 1, UnsupportedExprSnafu { @@ -4587,7 +4647,28 @@ impl PromPlanner { } ); - prev_field_exprs + self.ctx + .field_columns + .iter() + .map(|col| { + let value = DfExpr::Column(Column::from_name(col)); + // Normalize non `Float64` inputs the same way the label projection does, + // so both sides agree on the formatted value: `prom_float_to_string` + // formats exactly like Prometheus, while arrow's `Float64 -> Utf8` cast + // would render `1.0`. + let value = if Self::field_column_type(input_plan.schema(), col) + == Some(&ArrowDataType::Float64) + { + value + } else { + DfExpr::Cast(Cast::new(Box::new(value), ArrowDataType::Float64)) + }; + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![value], + }) + }) + .collect() } else { vec![] }; @@ -4969,7 +5050,7 @@ impl PromPlanner { .map(|col| { let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1); // Order by value in the specific order - sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true)); + sort_exprs.push(DfExpr::Column(Column::from_name(col)).sort(asc, true)); // Then tags if the values are equal, // Try to ensure the relative stability of the output results. sort_exprs.extend(tag_sort_exprs.clone()); @@ -4994,423 +5075,6 @@ impl PromPlanner { Ok(normalized_exprs) } - /// Create a classic, native, or mixed histogram helper plan. - async fn create_histogram_plan( - &mut self, - function_name: &str, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - let float_literal = |param: &PromExpr| -> Result { - let value = (|| { - let expr = Self::get_param_as_literal_expr( - Some(param), - None, - Some(ArrowDataType::Float64), - ) - .ok()?; - let simplifier = ExprSimplifier::new(SimplifyContext::default()); - let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; - let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { - return None; - }; - let ScalarValue::Float64(Some(value)) = - value.cast_to(&ArrowDataType::Float64).ok()? - else { - return None; - }; - Some(value) - })() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: function_name.to_string(), - })?; - Ok(value) - }; - let (function, input) = match (function_name, args.args.as_slice()) { - (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( - HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), - input.as_ref().clone(), - ), - (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( - HistogramFoldOperation::Fraction { - lower: float_literal(lower)?.into(), - upper: float_literal(upper)?.into(), - }, - input.as_ref().clone(), - ), - _ => { - return FunctionInvalidArgumentSnafu { - fn_name: function_name.to_string(), - } - .fail(); - } - }; - - let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; - // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a - // stable series identifier anymore. HistogramFold must not treat it as a label column. - let input_plan = self.strip_tsid_column(input_plan)?; - self.ctx.use_tsid = false; - - if let Some((float_field, histogram_field)) = - Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) - .map(|(float, histogram)| (float.to_string(), histogram.to_string())) - { - if self.ctx.has_le_tag() { - return self.create_mixed_histogram_plan( - function, - input_plan, - float_field, - histogram_field, - ); - } - self.ctx.field_columns = vec![histogram_field]; - } - if self.all_field_columns_are_native_histograms(input_plan.schema()) { - return self.create_native_histogram_plan(function, input_plan); - } - - if !self.ctx.has_le_tag() { - // Return empty result instead of error when 'le' column is not found - // This handles the case when histogram metrics don't exist - return Ok(LogicalPlan::EmptyRelation( - datafusion::logical_expr::EmptyRelation { - produce_one_row: false, - schema: input_plan.schema().clone(), - }, - )); - } - let time_index_column = - self.ctx - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - // FIXME(ruihang): support multi fields - let field_column = self - .ctx - .field_columns - .first() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: function.function_name().to_string(), - })? - .clone(); - // remove le column from tag columns - self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); - - let fold = HistogramFold::new_with_operation( - LE_COLUMN_NAME.to_string(), - field_column, - time_index_column, - function, - None, - input_plan, - ) - .context(DataFusionPlanningSnafu)?; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(fold), - })) - } - - fn create_native_histogram_expr( - &self, - function: HistogramFoldOperation, - field_column: &str, - ) -> DfExpr { - let field = DfExpr::Column(Column::from_name(field_column)); - let (func, args) = match function { - HistogramFoldOperation::Quantile(quantile) => ( - Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( - self.promql_annotations.clone(), - )), - vec![field, lit(f64::from(quantile))], - ), - HistogramFoldOperation::Fraction { lower, upper } => ( - Arc::new(NativeHistogramFraction::scalar_udf_with_collector( - self.promql_annotations.clone(), - )), - vec![field, lit(f64::from(lower)), lit(f64::from(upper))], - ), - }; - DfExpr::ScalarFunction(ScalarFunction { func, args }) - } - - fn create_native_histogram_plan( - &mut self, - function: HistogramFoldOperation, - input_plan: LogicalPlan, - ) -> Result { - ensure!( - self.ctx.field_columns.len() == 1, - MultiFieldsNotSupportedSnafu { - operator: function.function_name() - }, - ); - - let field_column = self.ctx.field_columns[0].clone(); - let function_expr = self.create_native_histogram_expr(function, &field_column); - let display_name = function_expr.schema_name().to_string(); - self.ctx.field_columns = vec![display_name.clone()]; - - let project_exprs = std::iter::once(self.create_time_index_column_expr()?) - .chain(std::iter::once(function_expr.alias(display_name))) - .chain(self.create_tag_column_exprs()?) - .collect::>(); - - LogicalPlanBuilder::from(input_plan) - .project(project_exprs) - .context(DataFusionPlanningSnafu)? - .filter(self.create_empty_values_filter_expr(false)?) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - fn create_mixed_histogram_plan( - &mut self, - function: HistogramFoldOperation, - input_plan: LogicalPlan, - float_field: String, - histogram_field: String, - ) -> Result { - let time_index_column = - self.ctx - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - let tag_columns = self.ctx.tag_columns.clone(); - let folded = HistogramFold::new_with_operation( - LE_COLUMN_NAME.to_string(), - float_field.clone(), - time_index_column.clone(), - function, - Some(histogram_field.clone()), - input_plan, - ) - .context(DataFusionPlanningSnafu)?; - let record_collision = DfExpr::ScalarFunction(ScalarFunction { - func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( - "vector contains a mix of classic and native histograms".to_string(), - self.promql_annotations.clone(), - )), - args: vec![col(&float_field), col(&histogram_field)], - }); - let keep = when( - col(&float_field) - .is_not_null() - .and(col(&histogram_field).is_not_null()), - record_collision, - ) - .otherwise(lit(true)) - .context(DataFusionPlanningSnafu)?; - - let native_expr = self.create_native_histogram_expr(function, &histogram_field); - let output_field = native_expr.schema_name().to_string(); - let value = DfExpr::ScalarFunction(ScalarFunction { - func: coalesce(), - args: vec![col(&float_field), native_expr], - }); - self.ctx.field_columns = vec![output_field.clone()]; - LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { - node: Arc::new(folded), - })) - .filter(keep) - .context(DataFusionPlanningSnafu)? - .project( - std::iter::once(col(&time_index_column)) - .chain(std::iter::once(value.alias(output_field))) - .chain(tag_columns.iter().map(col)), - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - /// Create a [SPECIAL_VECTOR_FUNCTION] plan - async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result { - if args.args.len() != 1 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_VECTOR_FUNCTION.to_string(), - } - .fail(); - } - let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; - - // reuse `SPECIAL_TIME_FUNCTION` as name of time index column - self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); - self.ctx.reset_table_name_and_schema(); - self.ctx.tag_columns = vec![]; - self.ctx.aggregation_field_labels.clear(); - self.ctx.field_columns = vec![greptime_value().to_string()]; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new( - EmptyMetric::new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - SPECIAL_TIME_FUNCTION.to_string(), - greptime_value().to_string(), - Some(lit), - ) - .context(DataFusionPlanningSnafu)?, - ), - })) - } - - /// Create a [SCALAR_FUNCTION] plan - async fn create_scalar_plan( - &mut self, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - ensure!( - args.len() == 1, - FunctionInvalidArgumentSnafu { - fn_name: SCALAR_FUNCTION - } - ); - let input = self - .prom_expr_to_plan(&args.args[0], query_engine_state) - .await?; - let input_schema = input.schema().clone(); - let alternative_samples = - Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns); - let histogram_fields = self - .ctx - .field_columns - .iter() - .filter(|field| Self::field_column_is_native_histogram(&input_schema, field)) - .count(); - ensure!( - self.ctx.field_columns.len() == 1 || alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: SCALAR_FUNCTION - }, - ); - let scalar_field = self - .ctx - .field_columns - .iter() - .find(|field| !Self::field_column_is_native_histogram(&input_schema, field)) - .or_else(|| self.ctx.field_columns.first()) - .cloned() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: SCALAR_FUNCTION, - })?; - let input = if histogram_fields == self.ctx.field_columns.len() { - // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN - // for every evaluation timestamp without attempting a Struct-to-Float64 cast. - LogicalPlanBuilder::from(input) - .filter(lit(false)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else if histogram_fields > 0 { - // A mixed vector contributes only its float samples to scalar(). - LogicalPlanBuilder::from(input) - .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else { - input - }; - let scalar_plan = LogicalPlan::Extension(Extension { - node: Arc::new( - ScalarCalculate::new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - input, - self.ctx.time_index_column.as_ref().unwrap(), - &self.ctx.tag_columns, - &scalar_field, - self.ctx.table_name.as_deref(), - ) - .context(PromqlPlanNodeSnafu)?, - ), - }); - // scalar plan have no tag columns - self.ctx.tag_columns.clear(); - self.ctx.aggregation_field_labels.clear(); - self.ctx.field_columns.clear(); - self.ctx - .field_columns - .push(scalar_plan.schema().field(1).name().clone()); - Ok(scalar_plan) - } - - /// Create a [SPECIAL_ABSENT_FUNCTION] plan - async fn create_absent_plan( - &mut self, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - if args.args.len() != 1 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_ABSENT_FUNCTION.to_string(), - } - .fail(); - } - let input = self - .prom_expr_to_plan(&args.args[0], query_engine_state) - .await?; - - let time_index_expr = self.create_time_index_column_expr()?; - let first_field_expr = - self.create_field_column_exprs()? - .pop() - .with_context(|| ValueNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - let first_value_expr = first_value(first_field_expr, vec![]); - - let ordered_aggregated_input = LogicalPlanBuilder::from(input) - .aggregate( - vec![time_index_expr.clone()], - vec![first_value_expr.clone()], - ) - .context(DataFusionPlanningSnafu)? - .sort(vec![time_index_expr.sort(true, false)]) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - let fake_labels = self - .ctx - .selector_matcher - .iter() - .filter_map(|matcher| match matcher.op { - MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())), - _ => None, - }) - .collect::>(); - - // Create the absent plan - let absent_plan = LogicalPlan::Extension(Extension { - node: Arc::new( - Absent::try_new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - self.ctx.time_index_column.as_ref().unwrap().clone(), - self.ctx.field_columns[0].clone(), - fake_labels, - ordered_aggregated_input, - ) - .context(DataFusionPlanningSnafu)?, - ), - }); - - // The absent series carries the equality matchers as labels, not the input's - // tags or value fields, so the input's field grouping labels no longer apply. - self.ctx.aggregation_field_labels.clear(); - Ok(absent_plan) - } - /// Try to build a DataFusion Literal Expression from PromQL Expr, return /// `None` if the input is not a literal expression. fn try_build_literal_expr(expr: &PromExpr) -> Option { @@ -5956,6 +5620,225 @@ impl PromPlanner { Ok((left_tag_columns, right_tag_columns, force_empty_join)) } + /// Result labels of a vector-vector binary operation, following Prometheus `resultMetric`: + /// `on(...)` keeps only the matching labels, `ignoring(...)` drops them, and a group modifier + /// keeps the "many" side's labels plus the `group_x(...)` labels taken from the "one" side. + /// + /// The flag of each entry tells which operand the label is projected from. `None` means the + /// operation keeps a whole operand tag set, which the default projection already does. + fn binary_result_labels( + left_context: &PromPlannerContext, + right_context: &PromPlannerContext, + modifier: &Option, + ) -> Option> { + let modifier = modifier.as_ref()?; + let (many_is_left, include) = match &modifier.card { + VectorMatchCardinality::OneToOne => (true, None), + VectorMatchCardinality::ManyToOne(labels) => (true, Some(labels)), + VectorMatchCardinality::OneToMany(labels) => (false, Some(labels)), + // Set operators keep their operands' labels and don't reach this path. + VectorMatchCardinality::ManyToMany => return None, + }; + + let Some(include) = include else { + let matching = modifier.matching.as_ref()?; + let reduced = |keep: bool, labels: &BTreeSet<&String>| { + left_context + .tag_columns + .iter() + .filter(|tag| labels.contains(tag) == keep) + .map(|tag| (true, tag.clone())) + .collect() + }; + return Some(match matching { + LabelModifier::Include(on) => reduced(true, &on.labels.iter().collect()), + LabelModifier::Exclude(ignoring) => { + reduced(false, &ignoring.labels.iter().collect()) + } + }); + }; + + let (many_context, one_context) = if many_is_left { + (left_context, right_context) + } else { + (right_context, left_context) + }; + let include = include.labels.iter().collect::>(); + // An included label the "one" side doesn't carry is deleted from the result, so it is + // dropped from the "many" side as well. + let mut labels = many_context + .tag_columns + .iter() + .filter(|tag| !include.contains(tag)) + .map(|tag| (many_is_left, tag.clone())) + .collect::>(); + labels.extend( + include + .into_iter() + .filter(|label| one_context.tag_columns.contains(label)) + .map(|label| (!many_is_left, label.clone())), + ); + Some(labels) + } + + /// Resolve [`Self::binary_result_labels`] against the join output. + fn binary_result_label_projection( + schema: &DFSchemaRef, + left_table_ref: &TableReference, + right_table_ref: &TableReference, + left_context: &PromPlannerContext, + right_context: &PromPlannerContext, + labels: Vec<(bool, String)>, + ) -> Result { + let mut exprs = Vec::with_capacity(labels.len()); + let mut names = Vec::with_capacity(labels.len()); + let mut aggregation_field_labels = Vec::new(); + let mut sources = HashSet::new(); + for (from_left, label) in labels { + let (table_ref, context) = if from_left { + (left_table_ref, left_context) + } else { + (right_table_ref, right_context) + }; + let field = schema + .qualified_field_with_name(Some(table_ref), &label) + .context(DataFusionPlanningSnafu)?; + exprs.push(DfExpr::Column(field.into())); + if context.aggregation_field_labels.contains(&label) { + aggregation_field_labels.push(label.clone()); + } + sources.insert(from_left); + names.push(label); + } + + // One operand contributing every one of its tags as the whole result label set keeps the + // one-to-one correspondence between its `__tsid` and a result series: no other operand + // value reaches the labels, and the matching gives each of its rows a single partner. + let source_context = match sources.into_iter().collect::>().as_slice() { + [true] => Some((left_table_ref, left_context)), + [false] => Some((right_table_ref, right_context)), + _ => None, + }; + let tsid = source_context + .filter(|(_, context)| { + context.use_tsid + && context.tag_columns.len() == names.len() + && context.tag_columns.iter().all(|tag| names.contains(tag)) + }) + .and_then(|(table_ref, _)| { + Self::optional_tsid_projection(schema, Some(table_ref), true) + }); + + Ok(BinaryResultLabels { + exprs, + names, + aggregation_field_labels, + tsid, + }) + } + + /// Whether the result of a binary operation can hold two series with the same labels, which + /// only [`Self::binary_result_labels`] can introduce: one-to-one matching on a subset of the + /// tags, or a group modifier that overwrites a label of the "many" side. + fn binary_result_labels_may_repeat( + left_context: &PromPlannerContext, + right_context: &PromPlannerContext, + modifier: &Option, + ) -> bool { + let Some(modifier) = modifier else { + return false; + }; + match &modifier.card { + VectorMatchCardinality::OneToOne => match &modifier.matching { + None => false, + Some(LabelModifier::Include(on)) => { + let on = on.labels.iter().collect::>(); + !left_context.tag_columns.iter().all(|tag| on.contains(tag)) + } + Some(LabelModifier::Exclude(ignoring)) => ignoring + .labels + .iter() + .any(|label| left_context.tag_columns.contains(label)), + }, + VectorMatchCardinality::ManyToOne(include) => include + .labels + .iter() + .any(|label| left_context.tag_columns.contains(label)), + VectorMatchCardinality::OneToMany(include) => include + .labels + .iter() + .any(|label| right_context.tag_columns.contains(label)), + VectorMatchCardinality::ManyToMany => false, + } + } + + /// Wrap `plan` in a check that fails the query when a match group holds more than one row at + /// a timestamp. `group_exprs` are the label columns of the group, resolved against `plan`. + fn assert_unique_match_group( + plan: LogicalPlan, + group_exprs: Vec, + group_labels: Vec, + time_index_expr: DfExpr, + violation: MatchGroupViolation, + ) -> Result { + let mut partition_by = group_exprs.clone(); + partition_by.push(time_index_expr); + // A label may carry the generated name, and the count column has to stay unambiguous + // against every column the operand already has. + let occupied_column_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(); + let mut next_suffix = 0; + let count_column = loop { + let name = match next_suffix { + 0 => MATCH_GROUP_COUNT_COLUMN.to_string(), + suffix => format!("{MATCH_GROUP_COUNT_COLUMN}_{suffix}"), + }; + next_suffix += 1; + if !occupied_column_names.contains(name.as_str()) { + break name; + } + }; + let count = DfExpr::WindowFunction(Box::new(WindowFunction { + fun: WindowFunctionDefinition::AggregateUDF(count_udaf()), + params: WindowFunctionParams { + args: vec![lit(1i64)], + partition_by, + order_by: vec![], + window_frame: WindowFrame::new(None), + null_treatment: None, + distinct: false, + filter: None, + }, + })) + .alias(count_column.as_str()); + + let output_exprs = plan + .schema() + .iter() + .map(|(qualifier, field)| DfExpr::Column(Column::new(qualifier.cloned(), field.name()))) + .collect::>(); + let assert_expr = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)), + args: std::iter::once(ident(count_column.as_str())) + .chain(group_exprs) + .collect(), + }); + + LogicalPlanBuilder::from(plan) + .window(vec![count]) + .context(DataFusionPlanningSnafu)? + .filter(assert_expr) + .context(DataFusionPlanningSnafu)? + .project(output_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + fn binary_modifier_preserves_tsid_join_key( &self, left_context: &PromPlannerContext, @@ -6023,7 +5906,7 @@ impl PromPlanner { && !force_empty_join && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]) && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]); - let (left, right) = if !only_join_time_index + let (left, right, left_matched_tags, right_matched_tags) = if !only_join_time_index && !use_tsid_join && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier) { @@ -6044,28 +5927,79 @@ impl PromPlanner { false, modifier, )?; - (left, right) + ( + left, + right, + aligned_left_context.tag_columns, + aligned_right_context.tag_columns, + ) } else { + ( + left, + right, + left_context.tag_columns.clone(), + right_context.tag_columns.clone(), + ) + }; + + // A join key that covers the whole tag set of a side already makes that side's match + // groups unique, and so does a join on `__tsid`. Only the reduced keys need the check. + let checked_matching = !only_join_time_index + && !force_empty_join + && !use_tsid_join + && !matches!( + modifier.as_ref().map(|modifier| &modifier.card), + Some(VectorMatchCardinality::ManyToMany) + ); + // The right operand is the "one" side of the matching, unless `group_right` swaps them. + let one_side_is_left = matches!( + modifier.as_ref().map(|modifier| &modifier.card), + Some(VectorMatchCardinality::OneToMany(_)) + ); + let (left, right) = if !checked_matching { (left, right) + } else if one_side_is_left { + ( + Self::assert_unique_one_side( + left, + &left_tag_columns, + &left_matched_tags, + left_time_index_column.as_deref(), + true, + )?, + right, + ) + } else { + ( + left, + Self::assert_unique_one_side( + right, + &right_tag_columns, + &right_matched_tags, + right_time_index_column.as_deref(), + false, + )?, + ) }; // push time index column if it exists - if let (Some(left_time_index_column), Some(right_time_index_column)) = - (left_time_index_column, right_time_index_column) - { + if let (Some(left_time_index_column), Some(right_time_index_column)) = ( + left_time_index_column.clone(), + right_time_index_column.clone(), + ) { left_tag_columns.insert(left_time_index_column); right_tag_columns.insert(right_time_index_column); } let right = LogicalPlanBuilder::from(right) - .alias(right_table_ref) + .alias(right_table_ref.clone()) .context(DataFusionPlanningSnafu)? .build() .context(DataFusionPlanningSnafu)?; // Inner Join on time index column to concat two operator - LogicalPlanBuilder::from(left) - .alias(left_table_ref) + let join_plan = LogicalPlanBuilder::from(left) + .alias(left_table_ref.clone()) .context(DataFusionPlanningSnafu)? .join_detailed( right, @@ -6085,7 +6019,90 @@ impl PromPlanner { ) .context(DataFusionPlanningSnafu)? .build() - .context(DataFusionPlanningSnafu) + .context(DataFusionPlanningSnafu)?; + + // The result labels no longer identify a series on their own: the "many" side can hold + // several series per result label set, which PromQL cannot represent. + let labels = checked_matching + .then(|| Self::binary_result_labels(left_context, right_context, modifier)) + .flatten() + .filter(|_| { + Self::binary_result_labels_may_repeat(left_context, right_context, modifier) + }); + let Some(labels) = labels else { + return Ok(join_plan); + }; + let schema = join_plan.schema().clone(); + let group_exprs = labels + .iter() + .map(|(from_left, label)| { + let table_ref = if *from_left { + &left_table_ref + } else { + &right_table_ref + }; + schema + .qualified_field_with_name(Some(table_ref), label) + .map(|field| DfExpr::Column(field.into())) + .context(DataFusionPlanningSnafu) + }) + .collect::>>()?; + let time_index_expr = left_time_index_column + .as_deref() + .map(|column| { + schema + .qualified_field_with_name(Some(&left_table_ref), column) + .map(|field| DfExpr::Column(field.into())) + .context(DataFusionPlanningSnafu) + }) + .transpose()? + .with_context(|| UnexpectedPlanExprSnafu { + desc: "vector matching on a plan without a time index", + })?; + let violation = if matches!( + modifier.as_ref().map(|modifier| &modifier.card), + Some(VectorMatchCardinality::OneToOne) + ) { + MatchGroupViolation::ImplicitManyToOne + } else { + MatchGroupViolation::AmbiguousGroupLabels + }; + Self::assert_unique_match_group( + join_plan, + group_exprs, + labels.into_iter().map(|(_, label)| label).collect(), + time_index_expr, + violation, + ) + } + + /// Guard the side of a vector matching that must hold one series per match group. + fn assert_unique_one_side( + plan: LogicalPlan, + join_keys: &BTreeSet, + tag_columns: &[String], + time_index_column: Option<&str>, + one_side_is_left: bool, + ) -> Result { + let Some(time_index_column) = time_index_column else { + return Ok(plan); + }; + if tag_columns.iter().all(|tag| join_keys.contains(tag)) { + return Ok(plan); + } + + let group_labels = join_keys.iter().cloned().collect::>(); + let group_exprs = group_labels + .iter() + .map(|label| DfExpr::Column(Column::from_name(label))) + .collect(); + Self::assert_unique_match_group( + plan, + group_exprs, + group_labels, + DfExpr::Column(Column::from_name(time_index_column)), + MatchGroupViolation::DuplicateOnOneSide { one_side_is_left }, + ) } fn selected_binary_match_labels( @@ -6259,203 +6276,6 @@ impl PromPlanner { matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row) } - /// Build a set operator (AND/OR/UNLESS) - fn set_op_on_non_field_columns( - &mut self, - mut left: LogicalPlan, - mut right: LogicalPlan, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - op: TokenType, - modifier: &Option, - ) -> Result { - let left_tag_col_set = left_context - .tag_columns - .iter() - .cloned() - .collect::>(); - let right_tag_col_set = right_context - .tag_columns - .iter() - .cloned() - .collect::>(); - - if matches!(op.id(), token::T_LOR) { - return self.or_operator( - left, - right, - left_tag_col_set, - right_tag_col_set, - left_context, - right_context, - modifier, - ); - } - - if let Some(modifier) = modifier { - ensure!( - matches!( - modifier.card, - VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany - ), - UnsupportedVectorMatchSnafu { - name: modifier.card.clone(), - }, - ); - } - - let output_context = left_context.clone(); - let visible_left_schema = left.schema().clone(); - let mut left_context = left_context; - let mut right_context = right_context; - let added_marker_to_left = if Self::only_temporality_match_label_mismatches( - &left_context, - &right_context, - modifier, - ) { - let aligned = Self::align_temporality_match_column( - left, - right, - &mut left_context, - &mut right_context, - )?; - left = aligned.0; - right = aligned.1; - aligned.2 - } else { - false - }; - - let mut left_tag_col_set = left_context - .tag_columns - .iter() - .cloned() - .collect::>(); - let mut right_tag_col_set = right_context - .tag_columns - .iter() - .cloned() - .collect::>(); - if let Some(matching) = modifier - .as_ref() - .and_then(|modifier| modifier.matching.as_ref()) - { - match matching { - LabelModifier::Include(on) => { - let mask = on.labels.iter().cloned().collect::>(); - left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect(); - right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect(); - } - LabelModifier::Exclude(ignoring) => { - for label in &ignoring.labels { - let _ = left_tag_col_set.remove(label); - let _ = right_tag_col_set.remove(label); - } - } - } - } - ensure!( - left_tag_col_set == right_tag_col_set, - CombineTableColumnMismatchSnafu { - left: left_tag_col_set.iter().cloned().collect::>(), - right: right_tag_col_set.iter().cloned().collect::>(), - } - ); - - let left_time_index = left_context.time_index_column.clone().unwrap(); - let right_time_index = right_context.time_index_column.clone().unwrap(); - - // alias right time index column if necessary - if left_context.time_index_column != right_context.time_index_column { - let right_project_exprs = right - .schema() - .fields() - .iter() - .map(|field| { - if field.name() == &right_time_index { - DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index) - } else { - DfExpr::Column(Column::from_name(field.name())) - } - }) - .collect::>(); - - right = LogicalPlanBuilder::from(right) - .project(right_project_exprs) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - } - - let join_keys = left_tag_col_set - .into_iter() - .chain([left_time_index]) - .collect::>(); - - ensure!( - left_context.field_columns.len() == 1 - || Self::field_columns_are_alternative_samples( - left.schema(), - &left_context.field_columns, - ), - MultiFieldsNotSupportedSnafu { - operator: "AND/UNLESS operator" - } - ); - // Generate join plan. - // All set operations in PromQL are "distinct" - let result = match op.id() { - token::T_LAND => LogicalPlanBuilder::from(left) - .distinct() - .context(DataFusionPlanningSnafu)? - .join_detailed( - right, - JoinType::LeftSemi, - (join_keys.clone(), join_keys), - None, - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu), - token::T_LUNLESS => LogicalPlanBuilder::from(left) - .distinct() - .context(DataFusionPlanningSnafu)? - .join_detailed( - right, - JoinType::LeftAnti, - (join_keys.clone(), join_keys), - None, - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu), - token::T_LOR => { - // OR is handled at the beginning of this function, as it cannot - // be expressed using JOIN like AND and UNLESS. - unreachable!() - } - _ => UnexpectedTokenSnafu { token: op }.fail(), - }?; - let result = if added_marker_to_left { - LogicalPlanBuilder::from(result) - .project(visible_left_schema.iter().map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - })) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else { - result - }; - - // AND/UNLESS preserve the complete left operand's visible columns and values; encoded - // markers are decoded. - self.ctx = output_context; - Ok(result) - } - fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> { match data_type { data_type if data_type.is_string() => Some(data_type), @@ -6511,601 +6331,6 @@ impl PromPlanner { } } - // TODO(ruihang): change function name - #[allow(clippy::too_many_arguments)] - fn or_operator( - &mut self, - left: LogicalPlan, - right: LogicalPlan, - left_tag_cols_set: HashSet, - right_tag_cols_set: HashSet, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> Result { - let left_is_empty = Self::is_zero_row_empty_relation(&left); - let right_is_empty = Self::is_zero_row_empty_relation(&right); - match (left_is_empty, right_is_empty) { - (true, false) => { - self.ctx = right_context; - return Ok(right); - } - (false, true) => { - self.ctx = left_context; - return Ok(left); - } - (true, true) => { - self.ctx = left_context; - return Ok(left); - } - (false, false) => {} - } - - ensure!( - !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), - UnexpectedPlanExprSnafu { - desc: "OR operator input has zero columns", - } - ); - let left_has_alternative_samples = - Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); - let right_has_alternative_samples = Self::field_columns_are_alternative_samples( - right.schema(), - &right_context.field_columns, - ); - ensure!( - left_context.field_columns.len() == 1 || left_has_alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: "OR operator" - } - ); - ensure!( - right_context.field_columns.len() == 1 || right_has_alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: "OR operator" - } - ); - - // prepare hash sets - let all_tags = left_tag_cols_set - .union(&right_tag_cols_set) - .cloned() - .collect::>(); - let left_qualifier = left.schema().qualified_field(0).0.cloned(); - let right_qualifier = right.schema().qualified_field(0).0.cloned(); - let left_qualifier_string = left_qualifier - .as_ref() - .map(|l| l.to_string()) - .unwrap_or_default(); - let right_qualifier_string = right_qualifier - .as_ref() - .map(|r| r.to_string()) - .unwrap_or_default(); - let left_time_index_column = - left_context - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let right_time_index_column = - right_context - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: right_qualifier_string.clone(), - })?; - let native_histogram_type = Self::native_histogram_arrow_type(); - let is_numeric = |data_type: &ArrowDataType| { - matches!( - data_type, - ArrowDataType::Int8 - | ArrowDataType::Int16 - | ArrowDataType::Int32 - | ArrowDataType::Int64 - | ArrowDataType::UInt8 - | ArrowDataType::UInt16 - | ArrowDataType::UInt32 - | ArrowDataType::UInt64 - | ArrowDataType::Float32 - | ArrowDataType::Float64 - ) - }; - let left_fields = left_context - .field_columns - .iter() - .map(|name| { - left.schema() - .iter() - .find(|(_, field)| field.name() == name) - .map(|(qualifier, field)| { - (name.clone(), qualifier.cloned(), field.data_type().clone()) - }) - .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) - }) - .collect::>>()?; - let right_fields = right_context - .field_columns - .iter() - .map(|name| { - right - .schema() - .iter() - .find(|(_, field)| field.name() == name) - .map(|(qualifier, field)| { - (name.clone(), qualifier.cloned(), field.data_type().clone()) - }) - .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) - }) - .collect::>>()?; - let left_field = &left_fields[0]; - let right_field = &right_fields[0]; - let left_field_col = &left_field.0; - let right_field_col = &right_field.0; - let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { - fields.iter().all(|(_, _, data_type)| { - is_numeric(data_type) || data_type == &native_histogram_type - }) - }; - let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { - if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR value fields have incompatible types: {:?} and {:?}", - left_fields - .iter() - .map(|(_, _, data_type)| data_type) - .collect::>(), - right_fields - .iter() - .map(|(_, _, data_type)| data_type) - .collect::>() - ), - } - .fail(); - } - true - } else { - (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) - || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) - }; - let target_field_type = if mixed_sample_types { - // Mixed vectors use the existing response representation: one nullable float column - // and one nullable native-histogram column. - ArrowDataType::Float64 - } else if left_field.2 == right_field.2 { - left_field.2.clone() - } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { - ArrowDataType::Float64 - } else { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR value fields have incompatible types: {:?} and {:?}", - left_field.2, right_field.2 - ), - } - .fail(); - }; - let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { - let mut reserved_names = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - for (name, _, _) in left_fields.iter().chain(&right_fields) { - reserved_names.remove(name); - } - reserved_names.extend(all_tags.iter().cloned()); - let unique_name = |prefix: &str, reserved_names: &mut HashSet| { - let mut index = 0; - loop { - let name = format!("{prefix}{index}"); - index += 1; - if reserved_names.insert(name.clone()) { - break name; - } - } - }; - let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); - let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); - (float_field, histogram_field) - } else { - (left_field_col.clone(), String::new()) - }; - let left_tag_types = left_tag_cols_set - .iter() - .map(|label| { - left.schema() - .fields() - .iter() - .find(|field| field.name() == label) - .map(|field| (label.clone(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) - }) - .collect::>>()?; - let right_tag_types = right_tag_cols_set - .iter() - .map(|label| { - right - .schema() - .fields() - .iter() - .find(|field| field.name() == label) - .map(|field| (label.clone(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) - }) - .collect::>>()?; - let mut target_tag_types = HashMap::with_capacity(all_tags.len()); - for label in &all_tags { - let Some(data_type) = - Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label)) - else { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR label {label} has incompatible types: {:?} and {:?}", - left_tag_types.get(label), - right_tag_types.get(label) - ), - } - .fail(); - }; - target_tag_types.insert(label.clone(), data_type); - } - let left_has_tsid = left - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); - let right_has_tsid = right - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); - - // step 0: fill all columns in output schema - let mut all_columns_set = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment - // (e.g. `unknown_metric or some_metric`). - if !(left_has_tsid && right_has_tsid) { - all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME); - } - // remove time index column - all_columns_set.remove(&left_time_index_column); - all_columns_set.remove(&right_time_index_column); - if mixed_sample_types { - for (name, _, _) in left_fields.iter().chain(&right_fields) { - all_columns_set.remove(name); - } - all_columns_set.extend(all_tags.iter().cloned()); - all_columns_set.insert(mixed_float_field_col.clone()); - all_columns_set.insert(mixed_histogram_field_col.clone()); - } else if left_field_col != right_field_col { - // remove field column in the right - all_columns_set.remove(right_field_col); - } - let mut all_columns = all_columns_set.into_iter().collect::>(); - // sort to ensure the generated schema is not volatile - all_columns.sort_unstable(); - // use left time index column name as the result time index column name - all_columns.insert(0, left_time_index_column.clone()); - let mut occupied_column_names = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - - // step 1: align schema using project, fill non-exist columns with null - let aligned_label_expr = |col: &String, source_types: &HashMap| { - let target_type = &target_tag_types[col]; - if let Some(source_type) = source_types.get(col) { - let expr = DfExpr::Column(Column::new(None::, col)); - if source_type == target_type { - expr - } else { - DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone()) - } - } else { - DfExpr::Literal( - Self::string_scalar_value(target_type, None) - .expect("target label type is a string"), - None, - ) - .alias(col.clone()) - } - }; - let null_histogram = - ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; - let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], - output_col: &String| { - if output_col == &mixed_float_field_col { - if let Some((name, qualifier, data_type)) = fields - .iter() - .find(|(_, _, data_type)| is_numeric(data_type)) - { - let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); - if data_type == &ArrowDataType::Float64 { - expr.alias(output_col) - } else { - DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64)) - .alias(output_col) - } - } else { - DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) - } - } else { - fields - .iter() - .find(|(_, _, data_type)| data_type == &native_histogram_type) - .map(|(name, qualifier, _)| { - DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) - }) - .unwrap_or_else(|| { - DfExpr::Literal(null_histogram.clone(), None).alias(output_col) - }) - } - }; - let left_proj_exprs = all_columns.iter().map(|col| { - if mixed_sample_types - && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) - { - mixed_value_expr(&left_fields, col) - } else if !mixed_sample_types - && col == left_field_col - && left_field.2 != target_field_type - { - DfExpr::Cast(Cast::new( - Box::new(DfExpr::Column(Column::new( - left_field.1.clone(), - left_field_col, - ))), - target_field_type.clone(), - )) - .alias(left_field_col.clone()) - } else if target_tag_types.contains_key(col) { - aligned_label_expr(col, &left_tag_types) - } else { - DfExpr::Column(Column::new(None::, col)) - } - }); - let right_time_index_expr = DfExpr::Column(Column::new( - right_qualifier.clone(), - right_time_index_column, - )) - .alias(left_time_index_column.clone()); - // The field column in right side may not have qualifier (it may be removed by join operation), - // so we need to find it from the schema. - // `skip（1)` to skip the time index column - let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { - // expr - if mixed_sample_types - && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) - { - mixed_value_expr(&right_fields, col) - } else if !mixed_sample_types && col == left_field_col { - let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); - if right_field.2 != target_field_type { - DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone())) - .alias(left_field_col.clone()) - } else if left_field_col != right_field_col { - expr.alias(left_field_col.clone()) - } else { - expr - } - } else if target_tag_types.contains_key(col) { - aligned_label_expr(col, &right_tag_types) - } else { - DfExpr::Column(Column::new(None::, col)) - } - }); - let right_proj_exprs = [right_time_index_expr] - .into_iter() - .chain(right_proj_exprs_without_time_index); - - let left_projected = LogicalPlanBuilder::from(left) - .project(left_proj_exprs) - .context(DataFusionPlanningSnafu)? - .alias(left_qualifier_string.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - let right_projected = LogicalPlanBuilder::from(right) - .project(right_proj_exprs) - .context(DataFusionPlanningSnafu)? - .alias(right_qualifier_string.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 2: compute match columns - let mut match_columns = if let Some(modifier) = modifier - && let Some(matching) = &modifier.matching - { - match matching { - // keeps columns mentioned in `on` - LabelModifier::Include(on) => on.labels.clone(), - // removes columns memtioned in `ignoring` - LabelModifier::Exclude(ignoring) => { - let ignoring = ignoring.labels.iter().cloned().collect::>(); - all_tags.difference(&ignoring).cloned().collect() - } - } - } else { - all_tags.iter().cloned().collect() - }; - // sort to ensure the generated plan is not volatile - match_columns.sort_unstable(); - match_columns.dedup(); - occupied_column_names.extend( - left_projected - .schema() - .fields() - .iter() - .chain(right_projected.schema().fields().iter()) - .map(|field| field.name().clone()), - ); - - let visible_schema = left_projected.schema().clone(); - let visible_left_exprs = left_projected - .schema() - .iter() - .map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - }) - .collect::>(); - let visible_right_exprs = right_projected - .schema() - .iter() - .map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - }) - .collect::>(); - let mut left_match_exprs = Vec::with_capacity(match_columns.len()); - let mut right_match_exprs = Vec::with_capacity(match_columns.len()); - let mut next_internal_column = 0; - - for label in &match_columns { - let left_field = if left_tag_cols_set.contains(label) { - Some( - left_projected - .schema() - .iter() - .find(|(_, field)| field.name() == label) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, - ) - } else { - None - }; - let right_field = if right_tag_cols_set.contains(label) { - Some( - right_projected - .schema() - .iter() - .find(|(_, field)| field.name() == label) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, - ) - } else { - None - }; - let data_type = match (left_field.as_ref(), right_field.as_ref()) { - (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { - left_type.clone() - } - (Some((_, left_type)), Some((_, right_type))) => { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" - ), - } - .fail(); - } - (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), - (None, None) => ArrowDataType::Utf8, - }; - let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else { - return UnexpectedPlanExprSnafu { - desc: format!("OR match label {label} must be a string"), - } - .fail(); - }; - let internal_name = loop { - let name = format!("__promql_or_match_{next_internal_column}"); - next_internal_column += 1; - if occupied_column_names.insert(name.clone()) { - break name; - } - }; - left_match_exprs.push(Self::normalized_match_key_expr( - label, - left_field, - &value_type, - &internal_name, - )); - right_match_exprs.push(Self::normalized_match_key_expr( - label, - right_field, - &value_type, - &internal_name, - )); - } - - let left_augmented = LogicalPlanBuilder::from(left_projected) - .project(visible_left_exprs.into_iter().chain(left_match_exprs)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - let right_augmented = LogicalPlanBuilder::from(right_projected) - .project(visible_right_exprs.into_iter().chain(right_match_exprs)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 3: build `UnionDistinctOn` with normalized internal match keys. - let visible_field_count = visible_schema.fields().len(); - let compare_key_indices = - (visible_field_count..visible_field_count + match_columns.len()).collect::>(); - let (time_qualifier, _) = visible_schema - .iter() - .find(|(_, field)| field.name() == &left_time_index_column) - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let ts_col_idx = left_augmented - .schema() - .iter() - .position(|(qualifier, field)| { - qualifier == time_qualifier && field.name() == &left_time_index_column - }) - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let union_distinct_on = UnionDistinctOn::try_new( - left_augmented, - right_augmented, - compare_key_indices, - ts_col_idx, - ) - .context(DataFusionPlanningSnafu)?; - let augmented_result = LogicalPlan::Extension(Extension { - node: Arc::new(union_distinct_on), - }); - let result = LogicalPlanBuilder::from(augmented_result) - .project(visible_schema.iter().map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - })) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 4: update context - let output_field_col = left_field_col.clone(); - let mut output_context = left_context; - let mut visible_tags = all_tags.into_iter().collect::>(); - visible_tags.sort_unstable(); - output_context.time_index_column = Some(left_time_index_column); - output_context.tag_columns = visible_tags; - output_context.field_columns = if mixed_sample_types { - vec![mixed_float_field_col, mixed_histogram_field_col] - } else { - vec![output_field_col] - }; - output_context.use_tsid = left_has_tsid && right_has_tsid; - self.ctx = output_context; - - Ok(result) - } - /// Build a projection that project and perform operation expr for every value columns. /// Non-value columns (tag and timestamp) will be preserved in the projection. /// @@ -7118,6 +6343,20 @@ impl PromPlanner { input: LogicalPlan, name_to_expr: F, ) -> Result + where + F: FnMut(&String) -> Result, + { + self.projection_for_each_field_column_with_labels(input, None, name_to_expr) + } + + /// Like [`Self::projection_for_each_field_column`], but projects `result_labels` instead of + /// the context tag columns when a binary operation derived its own result label set. + fn projection_for_each_field_column_with_labels( + &mut self, + input: LogicalPlan, + result_labels: Option<&BinaryResultLabels>, + name_to_expr: F, + ) -> Result where F: FnMut(&String) -> Result, { @@ -7128,22 +6367,30 @@ impl PromPlanner { let table_ref = self.ctx.table_name.clone().map(TableReference::bare); // Derived labels can be unqualified even when the context still names the source table. let input_schema = input.schema().clone(); - let non_field_columns_iter = self - .ctx - .tag_columns - .iter() - .chain(self.ctx.time_index_column.iter()) - .map(|col| { - input_schema - .qualified_field_with_name(table_ref.as_ref(), col) - .or_else(|_| input_schema.qualified_field_with_unqualified_name(col)) - .map(|field| DfExpr::Column(field.into())) - .context(DataFusionPlanningSnafu) - }); - let tsid_iter = - Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid) - .into_iter() - .map(Ok); + let lookup = |col: &String| { + input_schema + .qualified_field_with_name(table_ref.as_ref(), col) + .or_else(|_| input_schema.qualified_field_with_unqualified_name(col)) + .map(|field| DfExpr::Column(field.into())) + .context(DataFusionPlanningSnafu) + }; + let tag_columns_iter = match result_labels { + Some(labels) => labels.exprs.iter().cloned().map(Ok).collect::>(), + None => self.ctx.tag_columns.iter().map(lookup).collect::>(), + }; + let non_field_columns_iter = tag_columns_iter + .into_iter() + .chain(self.ctx.time_index_column.iter().map(lookup)); + let tsid_iter = match result_labels { + Some(labels) => labels.tsid_projection(table_ref.clone()), + None => Self::optional_tsid_projection( + input.schema(), + table_ref.as_ref(), + self.ctx.use_tsid, + ), + } + .into_iter() + .map(Ok); // build computation exprs let result_field_columns = self @@ -7215,7 +6462,7 @@ impl PromPlanner { /// Generate an expr like `date_part("hour", )`. Caller should ensure the /// time index column in context is set fn date_part_on_time_index(&self, date_part: &str) -> Result { - let input_expr = datafusion::logical_expr::col( + let input_expr = datafusion::logical_expr::ident( self.ctx .time_index_column .as_ref() @@ -7324,7235 +6571,4 @@ enum ScalarFunc { } #[cfg(test)] -mod test { - use std::time::{Duration, UNIX_EPOCH}; - - use catalog::RegisterTableRequest; - use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; - use common_base::Plugins; - use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; - use common_query::native_histogram::{ - CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, - }; - use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; - use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; - use common_query::test_util::DummyDecoder; - use common_recordbatch::RecordBatch as GreptimeRecordBatch; - use datafusion::arrow::array::{ - Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, - }; - use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; - use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::datasource::source::DataSourceExec; - use datafusion::datasource::{MemTable, provider_as_source}; - use datafusion::execution::context::SessionContext; - use datafusion::logical_expr::Extension; - use datatypes::prelude::ConcreteDataType; - use datatypes::schema::{ColumnSchema, Schema}; - use promql_parser::label::Labels; - use promql_parser::parser; - use session::context::QueryContext; - use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; - use table::Table; - use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; - use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; - - use super::*; - use crate::QueryEngineContext; - use crate::options::QueryOptions; - use crate::parser::QueryLanguageParser; - use crate::query_engine::DefaultSerializer; - - mod delta; - - fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { - if let LogicalPlan::Extension(Extension { node }) = plan - && let Some(instant_manipulate) = node.as_any().downcast_ref::() - { - return Some(instant_manipulate); - } - - plan.inputs().into_iter().find_map(find_instant_manipulate) - } - - fn build_query_engine_state() -> QueryEngineState { - QueryEngineState::new( - new_memory_catalog_manager().unwrap(), - None, - None, - None, - None, - None, - false, - Plugins::default(), - QueryOptions::default(), - ) - } - - #[test] - fn common_label_type_preserves_only_shared_dictionary_encoding() { - let dictionary = ArrowDataType::Dictionary( - Box::new(ArrowDataType::UInt32), - Box::new(ArrowDataType::Utf8), - ); - let other_dictionary = ArrowDataType::Dictionary( - Box::new(ArrowDataType::Int32), - Box::new(ArrowDataType::Utf8), - ); - - assert_eq!( - Some(dictionary.clone()), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), None) - ); - } - - async fn build_optimized_promql_plan( - table_provider: DfTableSourceProvider, - eval_stmt: &EvalStmt, - ) -> LogicalPlan { - let state = build_query_engine_state(); - let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state) - .await - .unwrap(); - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - state - .optimize_by_extension_rules(raw_plan, &context) - .unwrap() - } - - async fn build_optimized_tsid_plan( - query: &str, - num_tag: usize, - num_field: usize, - end_secs: u64, - lookback_secs: u64, - ) -> String { - let eval_stmt = EvalStmt { - expr: parser::parse(query).unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(end_secs)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(lookback_secs), - }; - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - num_tag, - num_field, - ) - .await; - - build_optimized_promql_plan(table_provider, &eval_stmt) - .await - .display_indent_schema() - .to_string() - } - - async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) { - let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; - - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0")); - assert!(plan_str.contains("Distinct:")); - assert!(plan_str.contains(expected_outer_agg), "{plan_str}"); - assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); - } - - async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) { - let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await; - assert!(!plan_str.contains("Distinct:"), "{plan_str}"); - } - - fn build_eval_stmt(expr: &str) -> EvalStmt { - EvalStmt { - expr: parser::parse(expr).unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - } - } - - enum DirectOrValue { - Float64(f64), - Int64(i64), - NativeHistogram(NativeHistogram), - Utf8(&'static str), - } - - impl DirectOrValue { - fn data_type(&self) -> ArrowDataType { - match self { - Self::Float64(_) => ArrowDataType::Float64, - Self::Int64(_) => ArrowDataType::Int64, - Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), - Self::Utf8(_) => ArrowDataType::Utf8, - } - } - fn array(&self) -> Arc { - match self { - Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), - Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), - Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), - Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), - } - } - } - - fn direct_or_histogram() -> NativeHistogram { - NativeHistogram { - schema: 0, - zero_threshold: 0.0, - sum: 1.0, - reset_hint: CounterResetHint::Unknown, - start_timestamp: None, - custom_values: vec![], - positive_spans: vec![], - negative_spans: vec![], - count: 1.0, - zero_count: 1.0, - positive_buckets: vec![], - negative_buckets: vec![], - } - } - - fn operator_metric_table( - name: &str, - table_id: u32, - tag: &str, - le: Option<&str>, - value: DirectOrValue, - ) -> table::TableRef { - let value_type = match &value { - DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), - DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), - DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), - DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), - }; - let tag_count = 1 + usize::from(le.is_some()); - let mut columns = vec![ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - if le.is_some() { - columns.push(ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.extend([ - ColumnSchema::new( - "ts".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new("v".to_string(), value_type, true), - ]); - let schema = Arc::new(Schema::new(columns)); - let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; - if let Some(le) = le { - arrays.push(Arc::new(StringArray::from(vec![le]))); - } - arrays.extend([ - Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, - value.array(), - ]); - let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - name, - GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), - table_id, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let value_index = tag_count + 1; - let meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![value_index]) - .next_column_id((value_index + 1) as u32) - .build() - .unwrap(); - let info = Arc::new( - TableInfoBuilder::default() - .table_id(table_id) - .name(name) - .meta(meta) - .build() - .unwrap(), - ); - Arc::new(Table::new( - info, - FilterPushDownType::Unsupported, - backing.data_source(), - )) - } - - fn operator_table_provider() -> DfTableSourceProvider { - let catalog = MemoryCatalogManager::with_default_setup(); - let tables = [ - operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), - operator_metric_table( - "lh", - 2_002, - "b", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), - operator_metric_table( - "rh", - 2_004, - "a", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), - operator_metric_table( - "bad_classic", - 2_006, - "d", - Some("broken"), - DirectOrValue::Float64(1.0), - ), - operator_metric_table( - "bad_native", - 2_007, - "d", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - ]; - for table in tables { - let info = table.table_info(); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: info.name.clone(), - table_id: info.ident.table_id, - table, - }) - .unwrap(); - } - DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn operator_eval_stmt(expr: &str) -> EvalStmt { - let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); - EvalStmt { - expr: parser::parse(expr).unwrap(), - start: time, - end: time, - interval: Duration::from_secs(1), - lookback_delta: Duration::from_secs(5), - } - } - - struct DirectOrSource { - name: &'static str, - empty: bool, - timestamp: i64, - tags: Vec<(&'static str, Option<&'static str>)>, - value: DirectOrValue, - } - - fn source( - name: &'static str, - empty: bool, - timestamp: i64, - tags: Vec<(&'static str, Option<&'static str>)>, - value: DirectOrValue, - ) -> DirectOrSource { - DirectOrSource { - name, - empty, - timestamp, - tags, - value, - } - } - - fn tagged_source( - name: &'static str, - empty: bool, - tag: (&'static str, Option<&'static str>), - value: DirectOrValue, - ) -> DirectOrSource { - source(name, empty, 1, vec![("job", Some("job")), tag], value) - } - - fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { - source(name, true, 1, vec![("job", Some("job"))], value) - } - - fn table(source: &DirectOrSource) -> Arc { - let mut fields = vec![Field::new( - "ts", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - )]; - fields.extend( - source - .tags - .iter() - .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), - ); - fields.push(Field::new("v", source.value.data_type(), true)); - let schema = Arc::new(ArrowSchema::new(fields)); - let partitions = if source.empty { - vec![vec![]] - } else { - let mut columns: Vec> = - vec![Arc::new(TimestampMillisecondArray::from(vec![ - source.timestamp, - ]))]; - columns.extend( - source - .tags - .iter() - .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), - ); - columns.push(source.value.array()); - vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] - }; - Arc::new(MemTable::try_new(schema, partitions).unwrap()) - } - - fn scan(source: &DirectOrSource) -> LogicalPlan { - LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) - .unwrap() - .build() - .unwrap() - } - - fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { - PromPlannerContext { - table_name: Some(qualifier.to_string()), - time_index_column: Some("ts".to_string()), - field_columns: vec![field.to_string()], - tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), - ..Default::default() - } - } - - fn or_modifier(expr: &str) -> Option { - let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { - unreachable!() - }; - expr.modifier - } - - async fn plan_direct_or( - left: LogicalPlan, - right: LogicalPlan, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> LogicalPlan { - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - planner - .or_operator( - left, - right, - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - modifier, - ) - .unwrap() - } - - async fn execute( - plan: LogicalPlan, - state: &QueryEngineState, - ) -> (LogicalPlan, Vec) { - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); - let physical = state - .session_state() - .create_physical_plan(&optimized) - .await - .unwrap(); - let batches = - datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) - .await - .unwrap(); - (optimized, batches) - } - - async fn run( - left: &DirectOrSource, - right: &DirectOrSource, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> (LogicalPlan, Vec) { - let plan = plan_direct_or( - scan(left), - scan(right), - left_context, - right_context, - modifier, - ) - .await; - execute(plan, &build_query_engine_state()).await - } - - async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { - let sample = |histogram: bool| { - if histogram { - DirectOrValue::NativeHistogram(direct_or_histogram()) - } else { - DirectOrValue::Float64(1.25) - } - }; - let left = tagged_source( - "lhs", - false, - ( - "k", - Some(if histogram_on_left { - "histogram" - } else { - "float" - }), - ), - sample(histogram_on_left), - ); - let right = tagged_source( - "rhs", - false, - ( - "k", - Some(if histogram_on_left { - "float" - } else { - "histogram" - }), - ), - sample(!histogram_on_left), - ); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left_context = direct_or_context("lhs", &["job", "k"], "v"); - let right_context = direct_or_context("rhs", &["job", "k"], "v"); - let plan = planner - .or_operator( - scan(&left), - scan(&right), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on(k) rhs"), - ) - .unwrap(); - (planner, plan) - } - - async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { - let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); - let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); - let row_count = histograms.len() + 1; - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "ts", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new("k", ArrowDataType::Utf8, false), - Field::new(&float_field, ArrowDataType::Float64, true), - Field::new( - &histogram_field, - native_histogram_value_type().as_arrow_type(), - true, - ), - ])); - let mut histogram_values = Vec::with_capacity(row_count); - histogram_values.push(None); - histogram_values.extend(histograms.into_iter().map(Some)); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), - Arc::new(StringArray::from_iter_values( - (0..row_count).map(|row| format!("kind_{row}")), - )), - Arc::new(Float64Array::from_iter( - (0..row_count).map(|row| (row == 0).then_some(1.25)), - )), - build_histogram_array(&histogram_values), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let planner = PromPlanner { - table_provider, - ctx: PromPlannerContext { - table_name: Some("mixed".to_string()), - time_index_column: Some("ts".to_string()), - field_columns: vec![float_field, histogram_field], - tag_columns: vec!["k".to_string()], - ..Default::default() - }, - promql_annotations: None, - }; - (planner, plan) - } - - fn assert_no_internal_or_keys(schema: &DFSchema) { - assert!( - schema - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")), - "{schema:?}" - ); - } - - fn values(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - batch - .column_by_name(column) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .collect() - } - - fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - let values = datafusion::arrow::compute::cast( - batch.column_by_name(column).unwrap(), - &ArrowDataType::Float64, - ) - .unwrap(); - values - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect::>() - }) - .collect() - } - - fn histograms(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - let values = batch - .column_by_name(column) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..values.len()).filter_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .collect() - } - - fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { - let mut rows = batches - .iter() - .flat_map(|batch| { - let values = batch - .column_by_name("v") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let labels = batch - .column_by_name("k") - .map(|column| column.as_any().downcast_ref::().unwrap()); - (0..batch.num_rows()).map(move |i| { - ( - values.value(i), - labels.and_then(|labels| { - (!labels.is_null(i)).then(|| labels.value(i).to_string()) - }), - ) - }) - }) - .collect::>(); - rows.sort_by(|left, right| left.0.total_cmp(&right.0)); - rows - } - - fn matrix_source( - name: &'static str, - k: Option>, - timestamp: i64, - value: f64, - ) -> DirectOrSource { - let mut tags = vec![("job", Some("job"))]; - if let Some(k) = k { - tags.push(("k", k)); - } - source(name, false, timestamp, tags, DirectOrValue::Float64(value)) - } - - fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { - direct_or_context( - name, - if k.is_some() { &["job", "k"] } else { &["job"] }, - "v", - ) - } - - async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "non_existent_histogram_bucket".to_string(), - ), - (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()), - ], - &["pod", "instance"], - ) - .await - } - - fn assert_normal_metric_schema(plan: &LogicalPlan) { - let fields = plan.schema().fields(); - assert_eq!(fields.len(), 4, "{fields:?}"); - assert!( - fields.iter().any(|field| field.name() == "pod"), - "{fields:?}" - ); - assert!( - fields.iter().any(|field| field.name() == "instance"), - "{fields:?}" - ); - assert!( - fields - .iter() - .any(|field| field.name() == greptime_timestamp()), - "{fields:?}" - ); - assert!( - fields.iter().any(|field| { - field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64 - }), - "{fields:?}" - ); - } - - async fn build_test_table_provider_with_distinct_tags( - table_tags: &[(&str, &[&str])], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (table_name, tags) in table_tags { - let mut columns = tags - .iter() - .map(|tag| { - ColumnSchema::new( - (*tag).to_string(), - ConcreteDataType::string_datatype(), - false, - ) - }) - .collect::>(); - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let table_meta = TableMetaBuilder::empty() - .schema(Arc::new(Schema::new(columns))) - .primary_key_indices((0..tags.len()).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name((*table_name).to_string()) - .meta(table_meta) - .build() - .unwrap(); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: (*table_name).to_string(), - table_id: 1024, - table: EmptyTable::from_table_info(&table_info), - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn contains_histogram_fold(plan: &LogicalPlan) -> bool { - matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::()) - || plan.inputs().into_iter().any(contains_histogram_fold) - } - - async fn build_set_op_context_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_distinct_tags(&[ - ("bucket_metric", &["job", "le"]), - ("normal_metric", &["job"]), - ("fallback_metric", &["instance"]), - ]) - .await - } - - async fn build_or_context_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_distinct_tags(&[ - ("normal_metric", &["job"]), - ("other_metric", &["instance"]), - ("non_hist_metric", &["instance"]), - ]) - .await - } - - async fn optimize_and_create_physical_plan( - state: &QueryEngineState, - plan: LogicalPlan, - ) -> ( - LogicalPlan, - Arc, - ) { - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); - let physical = state - .session_state() - .create_physical_plan(&optimized) - .await - .unwrap(); - (optimized, physical) - } - - async fn build_test_table_provider( - table_name_tuples: &[(String, String)], - num_tag: usize, - num_field: usize, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (schema_name, table_name) in table_name_tuples { - let mut columns = vec![]; - for i in 0..num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..num_tag).collect()) - .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id: 1024, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { - build_test_native_histogram_table_provider_with_marker(table_name, false).await - } - - async fn build_test_native_histogram_table_provider_with_marker( - table_name: &str, - temporality_marker: bool, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let mut columns = vec![ - ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - true, - ), - ]; - if temporality_marker { - columns.push(ColumnSchema::new( - OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), - ConcreteDataType::string_datatype(), - true, - )); - } - let tag_count = columns.len(); - columns.extend([ - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ]); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![tag_count + 1]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let columns = vec![ - ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - "native_histogram_2".to_string(), - native_histogram_value_type().clone(), - true, - ), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2, 3]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_mixed_native_histogram_table_provider( - table_name: &str, - ) -> DfTableSourceProvider { - build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await - } - - async fn build_test_mixed_native_histogram_table_provider_with_marker( - table_name: &str, - temporality_marker: bool, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let mut columns = vec![ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - if temporality_marker { - columns.push(ColumnSchema::new( - OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), - ConcreteDataType::string_datatype(), - true, - )); - } - let tag_count = columns.len(); - columns.extend([ - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ]); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![tag_count + 1, tag_count + 2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(), - ); - let mut arrays: Vec> = - vec![Arc::new(StringArray::from(vec!["float", "histogram"]))]; - if temporality_marker { - arrays.push(Arc::new(StringArray::from(vec![ - Some(GREPTIME_TEMPORALITY_DELTA), - Some(GREPTIME_TEMPORALITY_DELTA), - ]))); - } - arrays.extend([ - Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc, - build_histogram_array(&[None, Some(direct_or_histogram())]), - Arc::new(Float64Array::from(vec![Some(2.0), None])), - ]); - let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - table_name, - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 1024, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn classic_and_native_histogram_table_provider( - native_tag: &str, - native_le: Option<&str>, - native_histogram: NativeHistogram, - ) -> DfTableSourceProvider { - let table_name = "mixed_histogram"; - let catalog = MemoryCatalogManager::with_default_setup(); - let schema = Arc::new(Schema::new(vec![ - ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - true, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ])); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices(vec![0, 1]) - .value_indices(vec![3, 4]) - .next_column_id(5) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(), - ); - let batch = RecordBatch::try_new( - schema.arrow_schema().clone(), - vec![ - Arc::new(StringArray::from(vec![ - "classic", "classic", native_tag, "classic", "classic", native_tag, - ])), - Arc::new(StringArray::from(vec![ - Some("1"), - Some("+Inf"), - native_le, - Some("1"), - Some("+Inf"), - native_le, - ])), - Arc::new(TimestampMillisecondArray::from(vec![ - 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, - ])), - build_histogram_array(&[ - None, - None, - Some(native_histogram.clone()), - None, - None, - Some(native_histogram), - ]), - Arc::new(Float64Array::from(vec![ - Some(2.0), - Some(4.0), - None, - Some(2.0), - Some(4.0), - None, - ])), - ], - ) - .unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - table_name, - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 2_200, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 2_200, - table, - }) - .unwrap(); - - DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_table_provider_with_tsid( - table_name_tuples: &[(String, String)], - num_tag: usize, - num_field: usize, - ) -> DfTableSourceProvider { - let table_specs = table_name_tuples - .iter() - .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field)) - .collect::>(); - build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await - } - - async fn build_test_table_provider_with_tsid_fields( - table_specs: &[((String, String), usize)], - num_tag: usize, - ) -> DfTableSourceProvider { - let table_specs = table_specs - .iter() - .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field)) - .collect::>(); - build_test_table_provider_with_tsid_tag_fields(&table_specs).await - } - - async fn build_test_table_provider_with_tsid_tag_fields( - table_specs: &[((String, String), usize, usize)], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - - let physical_table_name = "phy"; - let physical_table_id = 999u32; - let physical_num_tag = table_specs - .iter() - .map(|(_, num_tag, _)| *num_tag) - .max() - .unwrap_or(0); - let physical_num_field = table_specs - .iter() - .map(|(_, _, num_field)| *num_field) - .max() - .unwrap_or(0); - - // Register a metric engine physical table with internal columns. - { - let mut columns = vec![ - ColumnSchema::new( - DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), - ConcreteDataType::uint32_datatype(), - false, - ), - ColumnSchema::new( - DATA_SCHEMA_TSID_COLUMN_NAME.to_string(), - ConcreteDataType::uint64_datatype(), - false, - ), - ]; - for i in 0..physical_num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..physical_num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - - let schema = Arc::new(Schema::new(columns)); - let primary_key_indices = (0..(2 + physical_num_tag)).collect::>(); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(primary_key_indices) - .value_indices( - (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(), - ) - .engine(METRIC_ENGINE_NAME.to_string()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(physical_table_id) - .name(physical_table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: physical_table_name.to_string(), - table_id: physical_table_id, - table, - }) - .is_ok() - ); - } - - // Register metric engine logical tables without `__tsid`, referencing the physical table. - for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate() - { - let mut columns = vec![]; - for i in 0..*num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..*num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - - let schema = Arc::new(Schema::new(columns)); - let mut options = table::requests::TableOptions::default(); - options.extra_options.insert( - LOGICAL_TABLE_METADATA_KEY.to_string(), - physical_table_name.to_string(), - ); - let table_id = 1024u32 + idx as u32; - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..*num_tag).collect()) - .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect()) - .engine(METRIC_ENGINE_NAME.to_string()) - .options(options) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(table_id) - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_table_provider_with_fields( - table_name_tuples: &[(String, String)], - tags: &[&str], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (schema_name, table_name) in table_name_tuples { - let mut columns = vec![]; - let num_tag = tags.len(); - for tag in tags { - columns.push(ColumnSchema::new( - tag.to_string(), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..num_tag).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id: 1024, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - // { - // input: `abs(some_metric{foo!="bar"})`, - // expected: &Call{ - // Func: MustGetFunction("abs"), - // Args: Expressions{ - // &VectorSelector{ - // Name: "some_metric", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // }, - // }, - // }, - async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) { - let prom_expr = - parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, 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]\ - \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]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]" - ).replace("TEMPLATE", plan_name); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn single_abs() { - do_single_instant_function_call("abs", "abs").await; - } - - #[tokio::test] - #[should_panic] - async fn single_absent() { - do_single_instant_function_call("absent", "").await; - } - - #[tokio::test] - async fn single_ceil() { - do_single_instant_function_call("ceil", "ceil").await; - } - - #[tokio::test] - async fn single_exp() { - do_single_instant_function_call("exp", "exp").await; - } - - #[tokio::test] - async fn single_ln() { - do_single_instant_function_call("ln", "ln").await; - } - - #[tokio::test] - async fn single_log2() { - do_single_instant_function_call("log2", "log2").await; - } - - #[tokio::test] - async fn single_log10() { - do_single_instant_function_call("log10", "log10").await; - } - - #[tokio::test] - #[should_panic] - async fn single_scalar() { - do_single_instant_function_call("scalar", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sgn() { - do_single_instant_function_call("sgn", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sort() { - do_single_instant_function_call("sort", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sort_desc() { - do_single_instant_function_call("sort_desc", "").await; - } - - #[tokio::test] - async fn single_sqrt() { - do_single_instant_function_call("sqrt", "sqrt").await; - } - - #[tokio::test] - async fn single_timestamp_plan_preserves_source_value() { - let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#); - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \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(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(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]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn single_acos() { - do_single_instant_function_call("acos", "acos").await; - } - - #[tokio::test] - #[should_panic] - async fn single_acosh() { - do_single_instant_function_call("acosh", "").await; - } - - #[tokio::test] - async fn single_asin() { - do_single_instant_function_call("asin", "asin").await; - } - - #[tokio::test] - #[should_panic] - async fn single_asinh() { - do_single_instant_function_call("asinh", "").await; - } - - #[tokio::test] - async fn single_atan() { - do_single_instant_function_call("atan", "atan").await; - } - - #[tokio::test] - #[should_panic] - async fn single_atanh() { - do_single_instant_function_call("atanh", "").await; - } - - #[tokio::test] - async fn single_cos() { - do_single_instant_function_call("cos", "cos").await; - } - - #[tokio::test] - #[should_panic] - async fn single_cosh() { - do_single_instant_function_call("cosh", "").await; - } - - #[tokio::test] - async fn single_sin() { - do_single_instant_function_call("sin", "sin").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sinh() { - do_single_instant_function_call("sinh", "").await; - } - - #[tokio::test] - async fn single_tan() { - do_single_instant_function_call("tan", "tan").await; - } - - #[tokio::test] - #[should_panic] - async fn single_tanh() { - do_single_instant_function_call("tanh", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_deg() { - do_single_instant_function_call("deg", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_rad() { - do_single_instant_function_call("rad", "").await; - } - - // { - // input: "avg by (foo)(some_metric)", - // expected: &AggregateExpr{ - // Op: AVG, - // Expr: &VectorSelector{ - // Name: "some_metric", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // PosRange: PositionRange{ - // Start: 13, - // End: 24, - // }, - // }, - // Grouping: []string{"foo"}, - // PosRange: PositionRange{ - // Start: 0, - // End: 25, - // }, - // }, - // }, - async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { - let prom_expr = parser::parse(&format!( - "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})", - )) - .unwrap(); - let mut eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // test group by - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 2, - 2, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected_no_without = String::from( - "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1: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, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); - assert_eq!( - plan.display_indent_schema().to_string(), - expected_no_without - ); - - // test group without - if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr { - *modifier = Some(LabelModifier::Exclude(Labels { - labels: vec![String::from("tag_1")].into_iter().collect(), - })); - } - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 2, - 2, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected_without = String::from( - "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1: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, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); - assert_eq!(plan.display_indent_schema().to_string(), expected_without); - } - - #[tokio::test] - async fn aggregate_sum() { - do_aggregate_expr_plan("sum", "sum").await; - } - - #[tokio::test] - async fn tsid_is_used_for_series_divide_when_available() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - assert!(plan_str.contains("__tsid ASC NULLS FIRST")); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - - let manipulate = find_instant_manipulate(&plan).unwrap(); - let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new( - 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")); - } - - #[tokio::test] - async fn default_binary_join_uses_tsid_when_available() { - let eval_stmt = build_eval_stmt("some_metric / some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!( - !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn reject_binary_fill_modifiers() { - let state = build_query_engine_state(); - - for query in [ - "some_metric + fill(0) some_alt_metric", - "some_metric + fill_left(0) some_alt_metric", - "some_metric + fill_right(0) some_alt_metric", - "(some_metric + fill(0) some_alt_metric) + some_metric", - ] { - let eval_stmt = build_eval_stmt(query); - let table_provider = build_test_table_provider(&[], 0, 0).await; - let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) - .await - .unwrap_err(); - - assert!( - matches!( - &err, - crate::promql::error::Error::UnsupportedExpr { name, .. } - if name == "PromQL fill modifiers" - ), - "{err}" - ); - } - } - - #[tokio::test] - async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { - for query in [ - "timestamp(some_metric) / some_metric", - "some_metric / timestamp(some_metric)", - ] { - let eval_stmt = build_eval_stmt(query); - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}"); - assert!( - plan_str.contains("lhs.tag_0 = rhs.tag_0"), - "{query}: {plan_str}" - ); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), - "{query}: {plan_str}" - ); - } - } - - #[tokio::test] - async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() { - let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job"); - - let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ - ( - (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()), - 2, - 1, - ), - ( - (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()), - 1, - 1, - ), - ]) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - - assert!( - plan_str.contains("Boolean(false)") || plan_str.contains("false"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn tsid_is_preserved_for_nested_default_binary_joins() { - let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_third_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn repeated_tsid_binary_operand_reuses_leaf_plan() { - let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn repeated_tsid_binary_operand_reuses_shorter_field_side() { - let eval_stmt = - build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - let value_columns = field_names - .iter() - .filter(|name| { - *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME - }) - .count(); - assert_eq!(value_columns, 1, "{field_names:?}"); - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1025)") - .count(), - 1, - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn binary_island_reuses_self_operand_without_join() { - let eval_stmt = build_eval_stmt("some_metric / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 1, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_reuses_leaf_across_two_branches() { - let eval_stmt = - build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_generated_alias_avoids_user_column_names() { - let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_fields( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - &["prom_v0", "__prom_v0"], - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan.schema().field_names(); - assert!(field_names.iter().any(|name| name.ends_with(".prom_v0"))); - assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0"))); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_clears_qualifier_for_nested_unary_projection() { - let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_keeps_distinct_matcher_leaves() { - let eval_stmt = build_eval_stmt( - "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}", - ); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_keeps_offset_leaves_distinct() { - let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_falls_back_for_group_modifier() { - let eval_stmt = build_eval_stmt( - "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric", - ); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_falls_back_for_comparison_filter() { - let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn tsid_binary_join_uses_shorter_field_side() { - let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - let value_columns = field_names - .iter() - .filter(|name| { - *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME - }) - .count(); - assert_eq!(value_columns, 1, "{field_names:?}"); - } - - #[tokio::test] - async fn comparison_binary_join_uses_shorter_field_side() { - let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - assert!( - field_names.iter().any(|name| name == "field_0"), - "{field_names:?}" - ); - assert!( - !field_names.iter().any(|name| name == "field_1"), - "{field_names:?}" - ); - } - - #[tokio::test] - async fn label_matching_modifier_disables_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn ignoring_absent_label_keeps_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() { - let eval_stmt = - build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_full_label_set_keeps_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_partial_label_set_disables_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ - ( - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - 2, - 1, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - 1, - 1, - ), - ]) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { - let eval_stmt = build_eval_stmt("some_metric > some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::from_eval_stmt(&eval_stmt), - promql_annotations: None, - }; - let plan = planner - .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), - "{plan_str}" - ); - assert!(planner.ctx.use_tsid, "{plan_str}"); - } - - #[tokio::test] - async fn comparison_bool_binary_join_uses_tsid_when_available() { - let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn scalar_count_count_range_keeps_full_window() { - let plan_str = build_optimized_tsid_plan( - "scalar(count(count(some_metric) by (tag_0)))", - 1, - 1, - 100_000, - 1, - ) - .await; - assert!(plan_str.contains("ScalarCalculate: tags=[]")); - assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]")); - assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]")); - } - - #[tokio::test] - async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() { - let plan_str = build_optimized_tsid_plan( - "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))", - 2, - 1, - 10, - 300, - ) - .await; - assert!(plan_str.contains("Distinct:"), "{plan_str}"); - } - - #[tokio::test] - async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() { - assert_nested_count_rewrite_applies( - "count(count(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]" - ) - .await; - } - - #[tokio::test] - async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() { - assert_nested_count_rewrite_applies( - "count(sum(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]" - ) - .await; - } - - #[tokio::test] - async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() { - for (query, expected_outer_agg) in [ - ( - "count(avg(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]", - ), - ( - "count(min(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]", - ), - ( - "count(max(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]", - ), - ( - "count(stddev(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]", - ), - ( - "count(stdvar(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]", - ), - ] { - assert_nested_count_rewrite_applies(query, expected_outer_agg).await; - } - } - - #[tokio::test] - async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() { - let count_plan = - build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1) - .await; - assert!( - !count_plan.contains("some_metric.field_0 IS NOT NULL"), - "{count_plan}" - ); - - for query in [ - "count(sum(some_metric) by (tag_0))", - "count(avg(some_metric) by (tag_0))", - "count(min(some_metric) by (tag_0))", - "count(max(some_metric) by (tag_0))", - "count(stddev(some_metric) by (tag_0))", - "count(stdvar(some_metric) by (tag_0))", - ] { - let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; - assert!( - plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"), - "{query}: {plan_str}" - ); - } - } - - #[tokio::test] - async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() { - assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await; - assert_nested_count_rewrite_missing( - "count(sum(irate(some_metric[1h])) by (tag_0))", - 2, - 300, - ) - .await; - } - - #[tokio::test] - async fn physical_table_name_is_not_leaked_in_plan() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("TableScan: phy"), "{plan}"); - assert!(plan_str.contains("SubqueryAlias: some_metric")); - assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)")); - assert!(!plan_str.contains("TableScan: some_metric")); - } - - #[tokio::test] - async fn sum_without_does_not_group_by_tsid() { - let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn topk_without_does_not_partition_by_tsid() { - let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let window_line = plan_str - .lines() - .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) - .unwrap(); - let partition_by = window_line - .split("PARTITION BY [") - .nth(1) - .and_then(|s| s.split("] ORDER BY").next()) - .unwrap(); - assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn sum_by_does_not_group_by_tsid() { - let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn aggregate_over_binary_time_function_expr() { - for op in ["sum", "min", "max", "avg"] { - let prom_expr = parser::parse(&format!( - "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)" - )) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(aggr_line.contains(op), "{plan_str}"); - assert!(aggr_line.contains("first_value"), "{plan_str}"); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME }) - ); - } - } - - #[tokio::test] - async fn topk_by_does_not_partition_by_tsid() { - let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let window_line = plan_str - .lines() - .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) - .unwrap(); - let partition_by = window_line - .split("PARTITION BY [") - .nth(1) - .and_then(|s| s.split("] ORDER BY").next()) - .unwrap(); - assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn selector_matcher_on_tsid_does_not_use_internal_column() { - let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet) { - if let LogicalPlan::Filter(filter) = plan { - datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap(); - } - for input in plan.inputs() { - collect_filter_cols(input, out); - } - } - - let mut filter_cols = HashSet::new(); - collect_filter_cols(&plan, &mut filter_cols); - assert!( - !filter_cols - .iter() - .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - } - - #[tokio::test] - async fn tsid_is_not_used_when_physical_table_is_missing() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let catalog_list = MemoryCatalogManager::with_default_setup(); - - // Register a metric engine logical table referencing a missing physical table. - let mut columns = vec![ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - "field_0".to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let schema = Arc::new(Schema::new(columns)); - let mut options = table::requests::TableOptions::default(); - options - .extra_options - .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string()); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .engine(METRIC_ENGINE_NAME.to_string()) - .options(options) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(1024) - .name("some_metric") - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "some_metric".to_string(), - table_id: 1024, - table, - }) - .unwrap(); - - let table_provider = DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ); - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); - assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - } - - #[tokio::test] - async fn tsid_is_carried_only_when_aggregate_preserves_label_set() { - let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("first_value") && plan_str.contains("__tsid")); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - - // Merging aggregate: label set is reduced, tsid should not be carried. - let prom_expr = parser::parse("sum(some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("first_value")); - } - - #[tokio::test] - async fn or_operator_with_unknown_metric_does_not_require_tsid() { - let prom_expr = parser::parse("unknown_metric or some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - } - - #[tokio::test] - async fn aggregate_avg() { - do_aggregate_expr_plan("avg", "avg").await; - } - - #[tokio::test] - #[should_panic] // output type doesn't match - async fn aggregate_count() { - do_aggregate_expr_plan("count", "count").await; - } - - #[tokio::test] - async fn aggregate_min() { - do_aggregate_expr_plan("min", "min").await; - } - - #[tokio::test] - async fn aggregate_max() { - do_aggregate_expr_plan("max", "max").await; - } - - #[tokio::test] - async fn aggregate_group() { - // Regression test for `group()` aggregator. - // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"})) - // should be plannable, and `group()` should produce constant 1 for each group. - let prom_expr = parser::parse( - "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "kubernetes_build_info".to_string(), - )], - &["cluster", "service", "job"], - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("max(Float64(1")); - } - - #[tokio::test] - async fn aggregate_stddev() { - do_aggregate_expr_plan("stddev", "stddev_pop").await; - } - - #[tokio::test] - async fn aggregate_stdvar() { - do_aggregate_expr_plan("stdvar", "var_pop").await; - } - - // TODO(ruihang): add range fn tests once exprs are ready. - - // { - // input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}", - // expected: &BinaryExpr{ - // Op: ADD, - // LHS: &VectorSelector{ - // Name: "a", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // RHS: &VectorSelector{ - // Name: "sum", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // VectorMatching: &VectorMatching{}, - // }, - // }, - #[tokio::test] - async fn binary_op_column_column() { - let prom_expr = - parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\ - \n Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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(\"foo\") AND 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]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - async fn indie_query_plan_compare>(query: &str, expected: T) { - let prom_expr = parser::parse(query).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - "greptime_private".to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref()); - } - - #[tokio::test] - async fn binary_op_literal_column() { - let query = r#"1 + some_metric{tag_0="bar"}"#; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn binary_op_literal_literal() { - let query = r#"1 + 1"#; - let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N] - TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn simple_bool_grammar() { - let query = "some_metric != bool 1.2345"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn bool_with_additional_arithmetic() { - let query = "some_metric + (1 == bool 2)"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn simple_unary() { - let query = "-some_metric"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn increase_aggr() { - let query = "increase(some_metric[5m])"; - let expected = String::from( - "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - async fn native_histogram_plan(query: &str) -> String { - let table_provider = build_test_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent_schema().to_string() - } - - #[tokio::test] - async fn native_histogram_count_uses_native_udf() { - let plan = native_histogram_plan("histogram_count(some_metric)").await; - - assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - assert!(!plan.contains("HistogramFold:"), "{plan}"); - } - - #[tokio::test] - async fn timestamp_filters_native_histogram_stale_marker_before_projection() { - let mut stale = direct_or_histogram(); - stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); - let table = operator_metric_table( - "stale_histogram", - 2_100, - "a", - None, - DirectOrValue::NativeHistogram(stale), - ); - let catalog = MemoryCatalogManager::with_default_setup(); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "stale_histogram".to_string(), - table_id: 2_100, - table, - }) - .unwrap(); - let provider = DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - provider, - &operator_eval_stmt("timestamp(stale_histogram)"), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - } - - #[tokio::test] - async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { - let histograms = build_histogram_array(&[None]); - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new( - greptime_native_histogram(), - histograms.data_type().clone(), - true, - ), - Field::new(greptime_value(), ArrowDataType::Float64, true), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1_000])), - histograms, - Arc::new(Float64Array::from(vec![f64::from_bits( - PROMETHEUS_STALE_NAN_BITS, - )])), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let input = LogicalPlan::Extension(Extension { - node: Arc::new(SeriesDivide::new( - Vec::new(), - "timestamp".to_string(), - input, - )), - }); - let input = LogicalPlan::Extension(Extension { - node: Arc::new(InstantManipulate::new( - 1_000, - 1_000, - 5_000, - 1_000, - 0, - "timestamp".to_string(), - Vec::new(), - Some(greptime_native_histogram().to_string()), - input, - )), - }); - // Match timestamp()'s parent projection, which otherwise prunes the companion lane. - let plan = LogicalPlanBuilder::from(input) - .project([col("timestamp")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - } - - #[tokio::test] - async fn native_histogram_rate_can_feed_count() { - let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; - - assert!(plan.contains("prom_native_histogram_rate"), "{plan}"); - assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - } - - #[tokio::test] - async fn native_histogram_quantile_skips_classic_fold() { - let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; - - assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); - assert!(!plan.contains("HistogramFold:"), "{plan}"); - assert!(plan.contains("some_metric.le"), "{plan}"); - // The phi literal is threaded into the native quantile UDF as its second argument. - assert!(plan.contains("Float64(0.9)"), "{plan}"); - // The empty-values filter drops NULL quantile results so the output is empty - // when all native histogram samples are dropped. - assert!(plan.contains("IS NOT NULL"), "{plan}"); - } - - #[tokio::test] - async fn mixed_native_histogram_quantile_uses_histogram_field() { - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_quantile(0.9, some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - - assert!( - plan.contains("prom_native_histogram_quantile(greptime_native_histogram"), - "{plan}" - ); - assert!(!plan.contains("EmptyRelation"), "{plan}"); - } - - #[tokio::test] - async fn mixed_histogram_helpers_execute_classic_and_native_samples() { - let state = build_query_engine_state(); - for (query, expected) in [ - ( - "histogram_quantile(0.5, mixed_histogram)", - vec![("classic", 1.0), ("native", 0.0)], - ), - ( - "histogram_fraction(-Inf, +Inf, mixed_histogram)", - vec![("classic", 1.0), ("native", 1.0)], - ), - ] { - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); - assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - let mut actual = batches - .iter() - .flat_map(|batch| { - let tags = batch - .column_by_name("tag") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column_by_name(&value_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) - }) - .collect::>(); - actual.sort_by_key(|(tag, _)| *tag); - assert_eq!(actual, expected, "{query}"); - } - } - - #[tokio::test] - async fn mixed_histogram_helpers_report_annotations() { - let state = build_query_engine_state(); - let mut native_histogram = direct_or_histogram(); - native_histogram.count = 2.0; - native_histogram.sum = f64::NAN; - for (native_tag, expected_rows, expected_warnings, expected_infos) in [ - ( - "classic", - 0, - vec!["vector contains a mix of classic and native histograms"], - vec![], - ), - ( - "native", - 2, - vec![], - vec!["input to histogram_quantile has NaN observations, result is skewed higher"], - ), - ] { - let collector = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - classic_and_native_histogram_table_provider( - native_tag, - None, - native_histogram.clone(), - ), - &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected_rows - ); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert_eq!(warnings, expected_warnings); - assert_eq!(infos, expected_infos); - } - } - - #[tokio::test] - async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { - let state = build_query_engine_state(); - let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); - stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider( - "classic", - Some("native"), - direct_or_histogram(), - ), - &stmt, - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert_eq!( - plan_text.matches("TableScan: mixed_histogram").count(), - 1, - "{plan_text}" - ); - - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let (_, batches) = execute(plan, &state).await; - let mut actual = batches - .iter() - .flat_map(|batch| { - let le = batch - .column_by_name(LE_COLUMN_NAME) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let timestamps = batch - .column_by_name("timestamp") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column_by_name(&value_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()).map(|row| { - ( - timestamps.value(row), - (!le.is_null(row)).then(|| le.value(row).to_string()), - values.value(row), - ) - }) - }) - .collect::>(); - actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); - assert_eq!( - actual, - vec![ - (1_000, None, 1.0), - (1_000, Some("native".to_string()), 0.0), - (2_000, None, 1.0), - (2_000, Some("native".to_string()), 0.0), - ] - ); - } - - #[tokio::test] - async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { - let state = build_query_engine_state(); - for native_le in [None, Some("native")] { - for query in [ - "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", - "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", - ] { - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider( - "native", - native_le, - direct_or_histogram(), - ), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - 0, - "native_le={native_le:?}, query={query}" - ); - } - } - } - - #[tokio::test] - async fn native_histogram_quantile_rejects_multi_field_input() { - let table_provider = build_test_multi_histogram_table_provider("some_metric").await; - let result = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_quantile(0.9, some_metric)"), - &build_query_engine_state(), - ) - .await; - - let err = result.expect_err("histogram_quantile on two native histogram fields must fail"); - assert!( - err.to_string() - .contains("Multi fields calculation is not supported in histogram_quantile"), - "{err}" - ); - } - - #[tokio::test] - async fn native_histogram_topk_uses_drop_udf() { - let plan = native_histogram_plan("topk(1, some_metric)").await; - - assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}"); - assert!( - plan.contains("Filter: prom_native_histogram_drop_float") - && plan.contains("IS NOT NULL"), - "{plan}" - ); - } - - #[tokio::test] - async fn mixed_or_topk_bottomk_ignore_native_histograms() { - for op in ["topk", "bottomk"] { - let collector = PromqlAnnotationCollector::default(); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - assert!( - plan.schema() - .fields() - .iter() - .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), - "{plan:?}" - ); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert_eq!( - infos, - vec![format!( - "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" - )] - ); - } - } - - #[tokio::test] - async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { - let plan = native_histogram_plan("scalar(some_metric)").await; - - assert!(plan.contains("ScalarCalculate"), "{plan}"); - assert!(plan.contains("Filter: Boolean(false)"), "{plan}"); - assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); - } - - #[tokio::test] - async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() { - for function in ["sort", "sort_desc"] { - let plan = native_histogram_plan(&format!("{function}(some_metric)")).await; - - assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); - assert!( - !plan.contains(&format!("Sort: {}", greptime_native_histogram())), - "{plan}" - ); - assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); - } - - for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] { - let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await; - - assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}"); - assert!(plan.contains(greptime_native_histogram()), "{plan}"); - assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); - } - } - - #[tokio::test] - async fn unsupported_native_histogram_functions_use_drop_udf() { - for query in [ - "deriv(some_metric[5m])", - "min_over_time(some_metric[5m])", - "quantile_over_time(0.9, some_metric[5m])", - "predict_linear(some_metric[5m], 60)", - "round(some_metric)", - "abs(some_metric)", - ] { - let plan = native_histogram_plan(query).await; - - assert!( - plan.contains("prom_native_histogram_drop_float"), - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn native_histogram_absent_over_time_uses_native_udf() { - let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await; - - assert!( - plan.contains("prom_native_histogram_absent_over_time"), - "{plan}" - ); - } - - #[tokio::test] - async fn native_histogram_all_function_arms_route_correctly() { - // Every native-histogram match arm in `create_function_expr` must route to the - // expected UDF when all field columns are native histograms. `holt_winters` shares - // the `double_exponential_smoothing` arm but is not registered in the promql - // parser (0.10), so it cannot be exercised through a query string. - let cases = [ - // Range functions routed to native histogram UDFs. - ( - "increase(some_metric[5m])", - "prom_native_histogram_increase", - ), - ("rate(some_metric[5m])", "prom_native_histogram_rate"), - ("delta(some_metric[5m])", "prom_native_histogram_delta"), - ("idelta(some_metric[5m])", "prom_native_histogram_idelta"), - ("irate(some_metric[5m])", "prom_native_histogram_irate"), - ("resets(some_metric[5m])", "prom_native_histogram_resets"), - ("changes(some_metric[5m])", "prom_native_histogram_changes"), - ( - "avg_over_time(some_metric[5m])", - "prom_native_histogram_avg_over_time", - ), - ( - "sum_over_time(some_metric[5m])", - "prom_native_histogram_sum_over_time", - ), - ( - "count_over_time(some_metric[5m])", - "prom_native_histogram_count_over_time", - ), - ( - "last_over_time(some_metric[5m])", - "prom_native_histogram_last_over_time", - ), - ( - "present_over_time(some_metric[5m])", - "prom_native_histogram_present_over_time", - ), - // Unsupported functions dropped with the float-null UDF. - ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"), - ( - "min_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "max_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "stddev_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "stdvar_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "quantile_over_time(0.9, some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "predict_linear(some_metric[5m], 60)", - "prom_native_histogram_drop_float", - ), - ( - "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", - "prom_native_histogram_drop_float", - ), - ("round(some_metric)", "prom_native_histogram_drop_float"), - ("rad(some_metric)", "prom_native_histogram_drop_float"), - ("deg(some_metric)", "prom_native_histogram_drop_float"), - ("sgn(some_metric)", "prom_native_histogram_drop_float"), - // Instant helper functions routed to native histogram UDFs. - ( - "histogram_count(some_metric)", - "prom_native_histogram_count", - ), - ("histogram_sum(some_metric)", "prom_native_histogram_sum"), - ("histogram_avg(some_metric)", "prom_native_histogram_avg"), - ( - "histogram_stddev(some_metric)", - "prom_native_histogram_stddev", - ), - ( - "histogram_stdvar(some_metric)", - "prom_native_histogram_stdvar", - ), - ( - "histogram_fraction(-2 + 1, 2 / 2, some_metric)", - "prom_native_histogram_fraction", - ), - ]; - - for (query, expected_udf) in cases { - let plan = native_histogram_plan(query).await; - assert!(plan.contains(expected_udf), "{query}\n{plan}"); - if query.starts_with("histogram_fraction") { - assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); - } - } - } - - #[tokio::test] - async fn mixed_native_histogram_ranges_use_coordinated_udfs() { - let dual_output = [ - "increase(some_metric[5m])", - "rate(some_metric[5m])", - "delta(some_metric[5m])", - "idelta(some_metric[5m])", - "irate(some_metric[5m])", - "avg_over_time(some_metric[5m])", - "sum_over_time(some_metric[5m])", - "last_over_time(some_metric[5m])", - ]; - let float_output = [ - "resets(some_metric[5m])", - "changes(some_metric[5m])", - "deriv(some_metric[5m])", - "min_over_time(some_metric[5m])", - "max_over_time(some_metric[5m])", - "count_over_time(some_metric[5m])", - "absent_over_time(some_metric[5m])", - "present_over_time(some_metric[5m])", - "stddev_over_time(some_metric[5m])", - "stdvar_over_time(some_metric[5m])", - "quantile_over_time(0.9, some_metric[5m])", - "predict_linear(some_metric[5m], 60)", - "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", - ]; - - for query in dual_output.iter().chain(float_output.iter()) { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); - assert_eq!( - plan.contains("prom_mixed_range_histogram"), - dual_output.contains(query), - "{query}\n{plan}" - ); - } - - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; - assert_eq!(plan, expected); - } - - #[tokio::test] - async fn mixed_native_histogram_rate_executes_real_ranges() { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new(greptime_value(), ArrowDataType::Float64, true), - Field::new( - greptime_native_histogram(), - native_histogram_value_type().as_arrow_type(), - true, - ), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), - Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), - build_histogram_array(&[None, Some(direct_or_histogram()), None]), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let collector = PromqlAnnotationCollector::default(); - let mut planner = PromPlanner { - table_provider: build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await, - ctx: PromPlannerContext { - start: 3000, - end: 3000, - interval: 1000, - range: Some(3000), - time_index_column: Some("timestamp".to_string()), - field_columns: vec![ - greptime_native_histogram().to_string(), - greptime_value().to_string(), - ], - ..Default::default() - }, - promql_annotations: Some(collector.clone()), - }; - let input = LogicalPlan::Extension(Extension { - node: Arc::new( - RangeManipulate::new( - 3000, - 3000, - 1000, - 0, - 3000, - "timestamp".to_string(), - planner.ctx.field_columns.clone(), - input, - ) - .unwrap(), - ), - }); - let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { - unreachable!() - }; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = Vec::new(); - collector.append_to(&mut warnings, &mut Vec::new()); - assert!( - warnings - .iter() - .any(|warning| warning.contains("mix of float and native histogram")) - ); - } - - #[tokio::test] - async fn native_histogram_mixed_field_table_behaves() { - // Exercise function planning after float and histogram samples have already been - // represented as alternative nullable fields. Histogram functions must select the - // histogram field without adding a NULL float field that would reject every row. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_count(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("prom_native_histogram_count"), - "{plan_str}" - ); - assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}"); - assert!( - plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"), - "{plan_str}" - ); - - // Value sorting keeps the float column and never sorts by the histogram column. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("sort(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("greptime_value ASC NULLS FIRST"), - "{plan_str}" - ); - assert!( - !plan_str.contains("greptime_native_histogram ASC"), - "{plan_str}" - ); - - // scalar() ignores histogram samples and evaluates only the float field. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("scalar(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("ScalarCalculate"), "{plan_str}"); - assert!( - plan_str.contains("greptime_value IS NOT NULL"), - "{plan_str}" - ); - - // Functions that preserve both alternative fields keep rows with either sample type. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - let filter = plan_str.lines().next().unwrap(); - assert!( - filter.starts_with("Filter: ") - && filter.contains("greptime_native_histogram IS NOT NULL") - && filter.contains(" OR ") - && filter.contains("greptime_value IS NOT NULL"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn less_filter_on_value() { - let query = "some_metric < 1.2345"; - let expected = String::from( - "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn count_over_time() { - let query = "count_over_time(some_metric[5m])"; - let expected = String::from( - "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - /// The outer `PromRangeManipulate` from a subquery must be preceded by - /// `Sort` + `PromSeriesDivide`. - #[tokio::test] - async fn count_over_time_subquery() { - let query = "count_over_time(some_metric[10m:1m])"; - let expected = String::from( - "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \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 PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn test_hash_join() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_sum".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["uri", "kubernetes_namespace", "kubernetes_pod_name"], - ) - .await; - // Should be ok - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ - \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ - \n SubqueryAlias: http_server_requests_seconds_sum\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_sum\ - \n SubqueryAlias: http_server_requests_seconds_count\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_count"; - assert_eq!(plan.to_string(), expected); - } - - #[tokio::test] - async fn test_nested_histogram_quantile() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "greptime_servers_grpc_requests_elapsed_bucket".to_string(), - )], - &["pod", "le", "path", "code", "container"], - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_histogram_quantile_binary_op() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // Arithmetic applied to a histogram_quantile() result. Regression for #8144: - // HistogramFold used to drop the input column qualifiers, so the binary-op - // projection failed to resolve the qualified tag column. - let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "http_request_duration_seconds_bucket".to_string(), - )], - &["pod", "le"], - ) - .await; - // Should plan without a "No field named ..." error. - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_parse_and_operator() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let cases = [ - r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, - r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, - ]; - - for case in cases { - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "kubelet_volume_stats_used_bytes".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "kubelet_volume_stats_capacity_bytes".to_string(), - ), - ], - &["namespace", "persistentvolumeclaim"], - ) - .await; - // Should be ok - let _ = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - } - - #[tokio::test] - async fn test_nested_binary_op() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) - - ( - sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m])) - or - vector(0) - )"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "nginx_ingress_controller_requests".to_string(), - )], - &["namespace", "job"], - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_parse_or_operator() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#" - sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) / - (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100) - or - 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) / - sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#; - - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) + - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 - or - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 - or - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#; - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) + - sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or - (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or - (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "background_waitevent_cnt".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "foreground_waitevent_cnt".to_string(), - ), - ], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#; - let table_provider = build_test_table_provider_with_fields( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "container_cpu_load_average_10s".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "container_spec_cpu_quota".to_string(), - ), - ], - &["cluster_name", "host_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn value_matcher() { - // template - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let cases = [ - // single equal matcher - ( - r#"some_metric{__field__="field_1"}"#, - vec![ - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // two equal matchers - ( - r#"some_metric{__field__="field_1", __field__="field_0"}"#, - vec![ - "some_metric.field_0", - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single not_eq matcher - ( - r#"some_metric{__field__!="field_1"}"#, - vec![ - "some_metric.field_0", - "some_metric.field_2", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // two not_eq matchers - ( - r#"some_metric{__field__!="field_1", __field__!="field_2"}"#, - vec![ - "some_metric.field_0", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // equal and not_eq matchers (no conflict) - ( - r#"some_metric{__field__="field_1", __field__!="field_0"}"#, - vec![ - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // equal and not_eq matchers (conflict) - ( - r#"some_metric{__field__="field_2", __field__!="field_2"}"#, - vec![ - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single regex eq matcher - ( - r#"some_metric{__field__=~"field_1|field_2"}"#, - vec![ - "some_metric.field_1", - "some_metric.field_2", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single regex not_eq matcher - ( - r#"some_metric{__field__!~"field_1|field_2"}"#, - vec![ - "some_metric.field_0", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - ]; - - for case in cases { - let prom_expr = parser::parse(case.0).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let mut fields = plan.schema().field_names(); - let mut expected = case.1.into_iter().map(String::from).collect::>(); - fields.sort(); - expected.sort(); - assert_eq!(fields, expected, "case: {:?}", case.0); - } - - let bad_cases = [ - r#"some_metric{__field__="nonexistent"}"#, - r#"some_metric{__field__!="nonexistent"}"#, - ]; - - for case in bad_cases { - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - assert!(plan.is_err(), "case: {:?}", case); - } - } - - #[tokio::test] - async fn custom_schema() { - let query = "some_alt_metric{__schema__=\"greptime_private\"}"; - let expected = String::from( - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - - let query = "some_alt_metric{__database__=\"greptime_private\"}"; - let expected = String::from( - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - - let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\ - \n Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn only_equals_is_supported_for_special_matcher() { - let queries = &[ - "some_alt_metric{__schema__!=\"greptime_private\"}", - "some_alt_metric{__schema__=~\"lalala\"}", - "some_alt_metric{__database__!=\"greptime_private\"}", - "some_alt_metric{__database__=~\"lalala\"}", - ]; - - for query in queries { - let prom_expr = parser::parse(query).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - "greptime_private".to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - assert!(plan.is_err(), "query: {:?}", query); - } - } - - #[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(); - let columns = vec![ - ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_nanosecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - "field".to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name("metrics".to_string()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "metrics".to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - let plan = PromPlanner::stmt_to_plan( - DfTableSourceProvider::new( - catalog_list.clone(), - false, - QueryContext::arc(), - DummyDecoder::arc(), - true, - ), - &EvalStmt { - expr: parser::parse("metrics{tag = \"1\"}").unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }, - &build_query_engine_state(), - ) - .await - .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(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( - catalog_list.clone(), - false, - QueryContext::arc(), - DummyDecoder::arc(), - true, - ), - &EvalStmt { - expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }, - &build_query_engine_state(), - ) - .await - .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(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]" - ); - } - - #[tokio::test] - async fn test_nonexistent_label() { - // template - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"some_metric{nonexistent="hi"}"#; - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_label_join() { - let prom_expr = parser::parse( - "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = - build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = r#" -Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] - Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; - - let ret = plan.display_indent_schema().to_string(); - assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); - } - - #[tokio::test] - async fn test_label_replace() { - let prom_expr = parser::parse( - "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = - build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = r#" -Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] - Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; - - let ret = plan.display_indent_schema().to_string(); - assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); - } - - #[tokio::test] - async fn label_replace_aggregation_queries_plan_successfully() { - let aggregate = - r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#; - let queries = [ - aggregate.to_string(), - format!("{aggregate} <= 10"), - format!("{aggregate} * 0.8"), - format!("0.8 * {aggregate}"), - format!("{aggregate} <= {aggregate} * 0.8"), - ]; - let state = build_query_engine_state(); - let mut failures = Vec::new(); - - for query in queries { - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - if let Err(error) = - PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await - { - failures.push(format!("{query}: {error:?}")); - } - } - - assert!(failures.is_empty(), "{}", failures.join("\n")); - } - - #[tokio::test] - async fn test_matchers_to_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = - r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - )], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]"; - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_topk_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\ - \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_count_values_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ - \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_value_alias() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series"); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = r#" -Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, greptime_timestamp:Timestamp(ms)] - Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] - Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; - assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); - } - - #[tokio::test] - async fn test_quantile_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_or_not_exists_table_label() { - let state = build_query_engine_state(); - let provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], - &["job"], - ) - .await; - let raw = PromPlanner::stmt_to_plan( - provider, - &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), - &state, - ) - .await - .unwrap(); - assert!( - raw.display_indent_schema() - .to_string() - .contains("__promql_or_match_0@") - ); - let (optimized, batches) = execute(raw, &state).await; - assert_no_internal_or_keys(optimized.schema()); - assert!(batches.iter().all(|batch| { - batch - .schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")) - })); - } - - #[tokio::test] - async fn test_histogram_quantile_missing_le_column() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // Test case: histogram_quantile with a table that doesn't have 'le' column - let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - - // Create a table provider with a table that doesn't have 'le' column - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "non_existent_histogram_bucket".to_string(), - )], - &["pod", "instance"], // Note: no 'le' column - ) - .await; - - // Should return empty result instead of error - let result = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - - // This should succeed now (returning empty result) instead of failing with "Cannot find column le" - assert!( - result.is_ok(), - "Expected successful plan creation with empty result, but got error: {:?}", - result.err() - ); - - // Verify that the result is an EmptyRelation - let plan = result.unwrap(); - match plan { - LogicalPlan::EmptyRelation(_) => { - // This is what we expect - } - _ => panic!("Expected EmptyRelation, but got: {:?}", plan), - } - } - - #[tokio::test] - async fn test_direct_or_normalizes_missing_match_labels() { - type Case<'a> = ( - Option>, - Option>, - i64, - i64, - &'a [(f64, Option<&'a str>)], - ); - - let modifier = or_modifier("lhs or on(k) rhs"); - #[rustfmt::skip] - let cases: &[Case<'_>] = &[ - (None, None, 1, 1, &[(1.0, None)]), - (None, Some(Some("")), 1, 1, &[(1.0, None)]), - (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), - (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), - (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), - (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), - (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), - (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), - ]; - for &(left, right, left_ts, right_ts, expected) in cases { - let (optimized, batches) = run( - &matrix_source("lhs", left, left_ts, 1.0), - &matrix_source("rhs", right, right_ts, 2.0), - matrix_context("lhs", left), - matrix_context("rhs", right), - &modifier, - ) - .await; - assert_no_internal_or_keys(optimized.schema()); - assert_eq!( - rows(&batches), - expected - .iter() - .map(|(value, label)| (*value, label.map(str::to_string))) - .collect::>() - ); - } - } - - #[tokio::test] - async fn test_direct_or_match_modifiers() { - for (modifier, left, right, expected) in [ - (None, "left", "right", 2), - (or_modifier("lhs or on(k) rhs"), "same", "same", 1), - (or_modifier("lhs or on() rhs"), "left", "right", 1), - (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), - ] { - let (_, batches) = run( - &matrix_source("lhs", Some(Some(left)), 1, 1.0), - &matrix_source("rhs", Some(Some(right)), 1, 2.0), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &modifier, - ) - .await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected - ); - } - } - - #[tokio::test] - async fn test_direct_or_nested_projection_uses_left_context() { - let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); - let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); - let raw = plan_direct_or( - scan(&left), - scan(&right), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - assert!(raw.schema().iter().any(|(qualifier, field)| { - qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" - })); - let nested = LogicalPlanBuilder::from(raw) - .project(vec![ - DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(DfExpr::Column(Column::new( - Some(TableReference::bare("lhs")), - "v", - ))), - op: Operator::Plus, - right: Box::new(lit(1.0)), - }) - .alias("v_plus"), - ]) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(nested, &build_query_engine_state()).await; - assert_eq!(values(&batches, "v_plus"), vec![2.0]); - } - - #[tokio::test] - async fn test_direct_or_skips_user_internal_key_name() { - const USER_TAG: &str = "__promql_or_match_0"; - let left = tagged_source( - "lhs", - false, - (USER_TAG, Some("left")), - DirectOrValue::Float64(1.0), - ); - let right = tagged_source( - "rhs", - false, - (USER_TAG, Some("right")), - DirectOrValue::Float64(2.0), - ); - let raw = plan_direct_or( - scan(&left), - scan(&right), - direct_or_context("lhs", &["job", USER_TAG], "v"), - direct_or_context("rhs", &["job", USER_TAG], "v"), - &or_modifier("lhs or on(missing_label) rhs"), - ) - .await; - assert!( - raw.display_indent_schema() - .to_string() - .contains("__promql_or_match_1@") - ); - let (_, batches) = execute(raw, &build_query_engine_state()).await; - assert!( - batches - .iter() - .all(|batch| batch.column_by_name(USER_TAG).is_some()) - ); - } - - #[tokio::test] - async fn test_direct_or_substrait_round_trip_with_normalized_key() { - let state = build_query_engine_state(); - let ctx = SessionContext::new_with_state(state.session_state()); - let catalog = Arc::new(MemoryCatalogProvider::new()); - catalog - .register_schema("public", Arc::new(MemorySchemaProvider::new())) - .unwrap(); - ctx.register_catalog("datafusion", catalog); - let left = matrix_source("lhs", Some(Some("")), 1, 1.0); - let right = matrix_source("rhs", None, 1, 2.0); - ctx.register_table( - TableReference::full("datafusion", "public", "lhs"), - table(&left), - ) - .unwrap(); - ctx.register_table( - TableReference::full("datafusion", "public", "rhs"), - table(&right), - ) - .unwrap(); - let raw = plan_direct_or( - ctx.table("datafusion.public.lhs") - .await - .unwrap() - .into_unoptimized_plan(), - ctx.table("datafusion.public.rhs") - .await - .unwrap() - .into_unoptimized_plan(), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - let decoded = DFLogicalSubstraitConvertor - .decode( - DFLogicalSubstraitConvertor - .encode(&raw, DefaultSerializer) - .unwrap(), - ctx.state(), - ) - .await - .unwrap(); - let (optimized, batches) = execute(decoded, &state).await; - assert_no_internal_or_keys(optimized.schema()); - assert!(batches.iter().all(|batch| { - batch - .schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")) - })); - assert_eq!(values(&batches, "v"), vec![1.0]); - } - - #[tokio::test] - async fn test_direct_or_numeric_value_types() { - let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); - let right = tagged_source( - "rhs", - false, - ("k", Some("rhs")), - DirectOrValue::Float64(0.5), - ); - let (optimized, batches) = run( - &left, - &right, - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - assert_eq!( - optimized - .schema() - .field_with_name(None, "v") - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!(values(&batches, "v"), vec![0.5]); - let provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider: provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left_context = direct_or_context("lhs", &["job"], "v"); - let right_context = direct_or_context("rhs", &["job"], "v"); - let error = planner - .or_operator( - scan(&job_source("lhs", DirectOrValue::Utf8("x"))), - scan(&job_source("rhs", DirectOrValue::Float64(1.0))), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on() rhs"), - ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("OR value fields have incompatible types") - ); - } - - #[tokio::test] - async fn test_or_with_histogram_quantile_missing_le_column() { - let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; - let eval_stmt = build_eval_stmt(case); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - assert_normal_metric_schema(&plan); - } - - #[tokio::test] - async fn test_or_with_right_empty_histogram_restores_left_context() { - let eval_stmt = build_eval_stmt( - r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#, - ); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_or_with_both_empty_histograms() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#, - ); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "left_histogram_bucket".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "right_histogram_bucket".to_string(), - ), - ], - &["pod", "instance"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - match plan { - LogicalPlan::EmptyRelation(relation) => { - assert!(!relation.produce_one_row); - assert!(!relation.schema.fields().is_empty()); - assert!( - relation - .schema - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64) - ); - assert!( - relation - .schema - .fields() - .iter() - .any(|field| field.name() == "pod") - ); - assert!( - !relation - .schema - .fields() - .iter() - .any(|field| field.name() == "instance") - ); - } - _ => panic!("Expected EmptyRelation, but got: {plan:?}"), - } - } - - #[tokio::test] - async fn test_nested_or_with_both_empty_histograms() { - for case in [ - r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#, - r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#, - ] { - let eval_stmt = build_eval_stmt(case); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "left_histogram_bucket".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "right_histogram_bucket".to_string(), - ), - ], - &["pod", "instance"], - ) - .await; - - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - } - - #[tokio::test] - async fn test_or_with_empty_histogram_modifiers() { - for case in [ - r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#, - r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#, - ] { - let eval_stmt = build_eval_stmt(case); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - assert_normal_metric_schema(&plan); - } - } - - #[tokio::test] - async fn test_unless_preserves_left_context_for_histogram() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#, - ); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &state, - ) - .await - .unwrap(); - assert!(contains_histogram_fold(&plan), "{plan:?}"); - let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await; - assert!(contains_histogram_fold(&optimized), "{optimized:?}"); - let batches = - datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) - .await - .unwrap(); - assert!(batches.iter().all(|batch| batch.num_rows() == 0)); - } - - #[tokio::test] - async fn test_and_preserves_left_context_for_histogram() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#, - ); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!(contains_histogram_fold(&plan), "{plan:?}"); - } - - #[tokio::test] - async fn test_and_preserves_left_context_when_le_is_missing() { - let eval_stmt = - build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}"); - assert!(!plan.schema().fields().is_empty()); - assert!(!contains_histogram_fold(&plan), "{plan:?}"); - } - - async fn build_matching_filter_plan(query: &str) -> String { - let table_provider = build_test_table_provider_with_distinct_tags(&[ - ("metric_a", &["host", "device"]), - ("metric_b", &["host", "device"]), - ]) - .await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent().to_string() - } - - /// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value - /// field that is neither a primary key nor a value column of the metric, so - /// `count by(status) (...)` still reports it among the aggregation's tag columns. - async fn build_test_table_provider_with_string_field( - table_tags: &[(&str, &[&str])], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (table_name, tags) in table_tags { - let mut columns = tags - .iter() - .map(|tag| { - ColumnSchema::new( - (*tag).to_string(), - ConcreteDataType::string_datatype(), - false, - ) - }) - .collect::>(); - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - columns.push(ColumnSchema::new( - "status".to_string(), - ConcreteDataType::string_datatype(), - true, - )); - let table_meta = TableMetaBuilder::empty() - .schema(Arc::new(Schema::new(columns))) - .primary_key_indices((0..tags.len()).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name((*table_name).to_string()) - .meta(table_meta) - .build() - .unwrap(); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: (*table_name).to_string(), - table_id: 1024, - table: EmptyTable::from_table_info(&table_info), - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_matching_filter_plan_with_string_field(query: &str) -> String { - let table_provider = build_test_table_provider_with_string_field(&[ - ("metric_a", &["host", "device"]), - ("metric_b", &["host", "device"]), - ]) - .await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent().to_string() - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_both_operands() { - for query in [ - r#"metric_a / metric_b{host="foo"}"#, - r#"metric_a / on(host, device) metric_b{host="foo"}"#, - r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#, - r#"metric_a / ignoring(device) metric_b{host="foo"}"#, - r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#, - r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!( - plan.matches(r#"host = Utf8("foo")"#).count(), - 2, - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn binary_matching_label_filter_skips_selecting_aggregations() { - // `topk` ranks its input, so filtering before it changes the candidate set. - let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#; - let plan = build_matching_filter_plan(query).await; - assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); - } - - #[tokio::test] - async fn binary_value_field_matcher_stays_on_its_own_operand() { - let value = greptime_value(); - for query in [ - format!(r#"metric_a / metric_b{{{value}="2"}}"#), - format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#), - ] { - let plan = build_matching_filter_plan(&query).await; - assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}"); - } - } - - #[tokio::test] - async fn binary_value_field_matcher_is_not_copied_across_aggregations() { - // `status` varies between the samples of one series, so filtering the other operand by it - // would drop the newest sample before sample selection (#9242). - let query = r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#; - let plan = build_matching_filter_plan_with_string_field(query).await; - assert_eq!( - plan.matches(r#"Utf8("ready")"#).count(), - 1, - "{query}\n{plan}" - ); - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() { - // `count`, not `sum`: the string value field is not summable. - let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#; - let plan = build_matching_filter_plan_with_string_field(query).await; - assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}"); - } - - #[tokio::test] - async fn binary_matching_label_filter_skips_unproven_expressions() { - for query in [ - r#"metric_a > on(host, device) metric_b{host="foo"}"#, - r#"metric_a / on(host) group_left metric_b{host="foo"}"#, - r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#, - r#"metric_a / on(device) metric_b{host="foo"}"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); - } - } - - #[tokio::test] - async fn test_or_context_uses_left_qualified_output() { - let case = r#"(normal_metric or other_metric) + 1"#; - let eval_stmt = build_eval_stmt(case); - let state = build_query_engine_state(); - let plan = - PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64), - "{plan:?}" - ); - let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await; - } - - #[tokio::test] - async fn test_or_context_uses_left_qualified_empty_histogram_output() { - let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#; - let eval_stmt = build_eval_stmt(case); - let plan = PromPlanner::stmt_to_plan( - build_or_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64), - "{plan:?}" - ); - } - - #[tokio::test] - async fn test_direct_or_preserves_float_and_native_histogram_samples() { - for histogram_on_left in [false, true] { - let (planner, plan) = mixed_direct_or(histogram_on_left).await; - - let float_field = &planner.ctx.field_columns[0]; - let histogram_field = &planner.ctx.field_columns[1]; - assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); - assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); - assert_eq!( - plan.schema() - .field_with_name(None, float_field) - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!( - plan.schema() - .field_with_name(None, histogram_field) - .unwrap() - .data_type(), - &native_histogram_value_type().as_arrow_type() - ); - - let (optimized, batches) = execute(plan, &build_query_engine_state()).await; - assert_no_internal_or_keys(optimized.schema()); - let mut sample_kinds = batches - .iter() - .flat_map(|batch| { - let values = batch.column_by_name(float_field).unwrap(); - let histograms = batch.column_by_name(histogram_field).unwrap(); - (0..batch.num_rows()) - .map(|row| (values.is_valid(row), histograms.is_valid(row))) - }) - .collect::>(); - sample_kinds.sort_unstable(); - assert_eq!(sample_kinds, vec![(false, true), (true, false)]); - } - } - - #[tokio::test] - async fn malformed_classic_bucket_does_not_drop_native_histogram() { - let state = build_query_engine_state(); - let collector = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - assert_eq!(values(&batches, &value_field), vec![0.0]); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert!(infos.is_empty()); - } - - #[tokio::test] - async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!( - plan_text.contains("prom_native_histogram_mul_scalar"), - "{plan_text}" - ); - assert!( - plan_text.contains("prom_native_histogram_scalar_mul"), - "{plan_text}" - ); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); - assert!(values(&batches, &float_field).is_empty()); - let mut sums = histograms(&batches, &histogram_field) - .into_iter() - .map(|histogram| histogram.sum) - .collect::>(); - sums.sort_by(f64::total_cmp); - assert_eq!(sums, vec![2.0, 3.0]); - } - - #[tokio::test] - async fn test_mixed_binary_operator_reports_only_dropped_samples() { - for (query, expected_rows, expected_infos) in [ - ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), - ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), - ("(lf or on(tag) lh) % on(tag) lh", 0, 1), - ] { - let state = build_query_engine_state(); - let annotations = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt(query), - &state, - Some(annotations.clone()), - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected_rows, - "{query}" - ); - let mut warnings = vec![]; - let mut infos = vec![]; - annotations.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty(), "{query}: {warnings:?}"); - assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); - } - } - - #[tokio::test] - async fn test_histogram_only_min_drops_empty_aggregate_group() { - // `min` over native-histogram-only input drops every sample in the group, so the - // NULL-valued aggregate row must be filtered out. Otherwise an outer expression - // like `group()` resurrects the group Prometheus considers unseen. - let state = build_query_engine_state(); - for query in ["min(lh)", "group(min(lh))"] { - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - 0, - "{query}" - ); - } - } - - #[tokio::test] - async fn test_mixed_min_drops_histogram_only_group() { - // With alternative float/histogram fields, `min by (tag)` keeps float-only groups - // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of - // emitting a NULL-valued row for them. - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - assert_eq!(values(&batches, &float_field), vec![2.0]); - } - - #[tokio::test] - async fn test_mixed_or_can_feed_another_or() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); - let mut float_values = values(&batches, &float_field); - float_values.sort_by(f64::total_cmp); - assert_eq!(float_values, vec![2.0, 7.0]); - assert_eq!(histograms(&batches, &histogram_field).len(), 1); - } - - #[tokio::test] - async fn test_mixed_fields_align_with_single_float_vector() { - let (planner, mixed) = mixed_direct_or(false).await; - let scale = tagged_source( - "scale", - false, - ("k", Some("float")), - DirectOrValue::Float64(2.0), - ); - let scale = scan(&scale); - let scale_fields = vec!["v".to_string()]; - let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { - unreachable!() - }; - - let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( - mixed.schema(), - scale.schema(), - &planner.ctx.field_columns, - &scale_fields, - binary.op, - false, - false, - ); - assert!(invalid_pairs.is_empty()); - assert_eq!( - groups - .iter() - .map(|(output, _)| output.clone()) - .collect::>(), - planner.ctx.field_columns - ); - assert_eq!(groups.len(), 2); - assert!( - groups - .iter() - .flat_map(|(_, pairs)| pairs) - .all(|(_, right)| *right == &scale_fields[0]) - ); - - let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( - scale.schema(), - mixed.schema(), - &scale_fields, - &planner.ctx.field_columns, - binary.op, - false, - false, - ); - assert!(invalid_pairs.is_empty()); - assert_eq!( - groups - .iter() - .map(|(output, _)| output.clone()) - .collect::>(), - planner.ctx.field_columns - ); - assert_eq!(groups.len(), 2); - assert!( - groups - .iter() - .flat_map(|(_, pairs)| pairs) - .all(|(left, _)| *left == &scale_fields[0]) - ); - } - - #[tokio::test] - async fn test_non_bool_comparison_filters_mixed_sample_lanes() { - let (planner, input) = mixed_direct_or(false).await; - let input_schema = input.schema().clone(); - let plan = planner - .filter_on_field_column(input, |field| { - if PromPlanner::field_column_is_native_histogram(&input_schema, field) { - Ok(lit(false)) - } else { - Ok(col(field).gt(lit(0.0))) - } - }) - .unwrap(); - let float_field = planner.ctx.field_columns[0].clone(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_field), vec![1.25]); - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - } - - #[tokio::test] - async fn test_mixed_left_and_unless_preserve_sample_lanes() { - for (expression, expected_sample_kind) in [ - ("lhs and on(k) mask", (false, true)), - ("lhs unless on(k) mask", (true, false)), - ] { - let (mut planner, left) = mixed_direct_or(false).await; - let left_context = planner.ctx.clone(); - let float_field = left_context.field_columns[0].clone(); - let histogram_field = left_context.field_columns[1].clone(); - let mask = tagged_source( - "mask", - false, - ("k", Some("histogram")), - DirectOrValue::Float64(1.0), - ); - let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { - unreachable!() - }; - let plan = planner - .set_op_on_non_field_columns( - left, - scan(&mask), - left_context, - direct_or_context("mask", &["job", "k"], "v"), - binary.op, - &binary.modifier, - ) - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - let sample_kinds = batches - .iter() - .flat_map(|batch| { - let floats = batch.column_by_name(&float_field).unwrap(); - let histograms = batch.column_by_name(&histogram_field).unwrap(); - (0..batch.num_rows()) - .map(|row| (floats.is_valid(row), histograms.is_valid(row))) - }) - .collect::>(); - assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); - } - } - - #[tokio::test] - async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt("some_metric * scalar(vector(2))"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let schema = plan.schema(); - assert_eq!( - schema - .field_with_unqualified_name(greptime_value()) - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!( - schema - .field_with_unqualified_name(greptime_native_histogram()) - .unwrap() - .data_type(), - &native_histogram_value_type().as_arrow_type() - ); - assert!( - plan.display_indent_schema() - .to_string() - .contains("prom_native_histogram_mul_scalar"), - "{plan:?}" - ); - } - - #[tokio::test] - async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(values(&batches, &float_field), vec![0.0]); - assert_eq!(histograms(&batches, &histogram_field).len(), 1); - } - - #[tokio::test] - async fn test_unary_negates_mixed_float_and_native_histogram_samples() { - for histogram_on_left in [false, true] { - let (mut planner, input) = mixed_direct_or(histogram_on_left).await; - let plan = planner.negate_field_columns(input).unwrap(); - assert!(PromPlanner::field_columns_are_alternative_samples( - plan.schema(), - &planner.ctx.field_columns - )); - let float_field = planner - .ctx - .field_columns - .iter() - .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap(); - let histogram_field = planner - .ctx - .field_columns - .iter() - .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, float_field), vec![-1.25]); - let histogram = batches - .iter() - .find_map(|batch| { - let values = batch - .column_by_name(histogram_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..values.len()).find_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .unwrap(); - assert_eq!(histogram.count, -1.0); - assert_eq!(histogram.sum, -1.0); - assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); - } - } - - #[tokio::test] - async fn test_native_histogram_sum_and_avg_execute_real_batches() { - for op_name in ["sum", "avg"] { - for incompatible in [false, true] { - let mut second = direct_or_histogram(); - if incompatible { - second.schema = CUSTOM_BUCKETS_SCHEMA; - second.custom_values = vec![1.0]; - } - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = - mixed_aggregate_input(vec![direct_or_histogram(), second]).await; - planner.promql_annotations = Some(collector.clone()); - let histogram_column = planner.ctx.field_columns[1].clone(); - planner.ctx.field_columns = vec![histogram_column.clone()]; - let input = LogicalPlanBuilder::from(input) - .project([col("ts"), col(&histogram_column)]) - .unwrap() - .build() - .unwrap(); - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse(&format!("{op_name}(mixed)")).unwrap() - else { - unreachable!() - }; - let (aggregate_exprs, _) = - planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter(planner.create_empty_values_filter_expr(false).unwrap()) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(infos.is_empty()); - if incompatible { - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - assert!(warnings.iter().any(|warning| { - warning - == &format!( - "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" - ) - })); - } else { - let histograms = histograms(&batches, &histogram_column); - assert_eq!(histograms.len(), 1); - let expected = if op_name == "sum" { 2.0 } else { 1.0 }; - assert_eq!(histograms[0].count, expected); - assert_eq!(histograms[0].sum, expected); - assert!(warnings.is_empty()); - } - } - } - } - - #[tokio::test] - async fn test_canonical_mixed_count_group_and_count_values_execute() { - let state = build_query_engine_state(); - for (query, expected) in [ - ("count(some_metric)", vec![2.0]), - ("group(some_metric)", vec![1.0]), - (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), - ] { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_sample_count")), - "{query}: {plan:?}" - ); - let value_fields = plan - .schema() - .fields() - .iter() - .filter(|field| { - matches!( - field.data_type(), - ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 - ) || field.data_type() == &native_histogram_value_type().as_arrow_type() - }) - .collect::>(); - assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); - assert_ne!( - value_fields[0].data_type(), - &native_histogram_value_type().as_arrow_type(), - "{query}: {plan:?}" - ); - let value_column = value_fields[0].name().clone(); - - let (_, batches) = execute(plan, &state).await; - let mut actual = numeric_values(&batches, &value_column); - actual.sort_by(f64::total_cmp); - assert_eq!(actual, expected, "{query}"); - - if query.starts_with("count_values") { - let mut sample_labels = batches - .iter() - .flat_map(|batch| { - batch - .column_by_name("sample") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .map(str::to_string) - }) - .collect::>(); - sample_labels.sort(); - let mut expected_labels = - vec!["2".to_string(), direct_or_histogram().promql_string()]; - expected_labels.sort(); - assert_eq!(sample_labels, expected_labels); - } - } - } - - #[tokio::test] - async fn test_mixed_or_sum_aggregates_each_sample_type() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("sum(lhs)").unwrap() - else { - unreachable!() - }; - - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_direct_or(false).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts"), col("k")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([ - col(&float_column), - col(&histogram_column), - col("ts"), - col("k"), - ]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_column), vec![1.25]); - let histogram = batches - .iter() - .find_map(|batch| { - let values = batch - .column_by_name(&histogram_column)? - .as_any() - .downcast_ref::()?; - (0..values.len()).find_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .unwrap(); - assert_eq!(histogram.count, 1.0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_direct_or(false).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col(&histogram_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert_eq!( - warnings, - vec![ - "sum: dropped aggregation result containing both float and native histogram samples" - ] - ); - } - - #[tokio::test] - async fn test_mixed_or_sum_drops_incompatible_mixed_group() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("sum(lhs)").unwrap() - else { - unreachable!() - }; - let mut custom = direct_or_histogram(); - custom.schema = CUSTOM_BUCKETS_SCHEMA; - custom.custom_values = vec![1.0]; - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col(&histogram_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.iter().any(|warning| { - warning - == "sum: dropped aggregation result containing both float and native histogram samples" - })); - } - - #[tokio::test] - async fn test_mixed_or_min_records_only_present_histograms() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("min(lhs)").unwrap() - else { - unreachable!() - }; - let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; - - for (histograms, expected_infos) in [ - (vec![], vec![]), - (vec![direct_or_histogram()], vec![expected_info]), - ] { - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_aggregate_input(histograms).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_ignored_histogram_filter_expr(op, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_column), vec![1.25]); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert_eq!(infos, expected_infos); - } - } - - #[tokio::test] - async fn test_mixed_or_value_aliases_do_not_replace_labels() { - let left = source( - "lhs", - false, - 1, - vec![("job", Some("job")), ("k", Some("float"))], - DirectOrValue::Float64(1.0), - ); - let right = source( - "rhs", - false, - 1, - vec![ - ("job", Some("job")), - ("k", Some("histogram")), - (greptime_value(), Some("value-label")), - ], - DirectOrValue::NativeHistogram(direct_or_histogram()), - ); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left = LogicalPlanBuilder::from(scan(&left)) - .project(vec![ - col("ts"), - col("job"), - col("k"), - col("v").alias(greptime_value()), - ]) - .unwrap() - .build() - .unwrap(); - let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); - let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); - let plan = planner - .or_operator( - left, - scan(&right), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on(k) rhs"), - ) - .unwrap(); - - assert_eq!( - plan.schema() - .field_with_name(None, greptime_value()) - .unwrap() - .data_type(), - &ArrowDataType::Utf8 - ); - assert!( - planner - .ctx - .field_columns - .iter() - .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) - ); - assert!(PromPlanner::field_columns_are_alternative_samples( - plan.schema(), - &planner.ctx.field_columns - )); - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); - let labels = batches - .iter() - .flat_map(|batch| { - batch - .column_by_name(greptime_value()) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .collect::>(); - assert_eq!(labels, vec!["value-label"]); - } - - #[tokio::test] - async fn test_mixed_or_routes_float_histogram_and_label_functions() { - for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { - let (mut planner, input) = mixed_direct_or(false).await; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { - unreachable!() - }; - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - exprs.extend(planner.create_tag_column_exprs().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - let values = batches - .iter() - .flat_map(|batch| { - batch - .schema() - .fields() - .iter() - .position(|field| field.data_type() == &ArrowDataType::Float64) - .map(|index| { - batch - .column(index) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .into_iter() - .flatten() - }) - .collect::>(); - assert_eq!(values, vec![expected], "{function}"); - } - - let (mut planner, input) = mixed_direct_or(false).await; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let PromExpr::Call(call) = - parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() - else { - unreachable!() - }; - let args = planner.create_function_args(&call.args.args).unwrap(); - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, args.literals, input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - exprs.extend(planner.create_tag_column_exprs().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); - assert_eq!(sample_count, 2); - } -} +mod test; diff --git a/src/query/src/promql/planner/at_modifier.rs b/src/query/src/promql/planner/at_modifier.rs new file mode 100644 index 00000000000..545f1714a22 --- /dev/null +++ b/src/query/src/promql/planner/at_modifier.rs @@ -0,0 +1,326 @@ +// 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. + +//! Anchored selector planning and replay for the PromQL `@` modifier. + +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr}; +use promql::extension_plan::{InstantManipulate, Millisecond, SeriesDivide}; +use promql_parser::parser::{ + AtModifier, Call, Expr as PromExpr, MatrixSelector, Offset, ParenExpr, +}; +use snafu::{OptionExt, ResultExt, ensure}; + +use crate::promql::error::{ + AtModifierTimestampOutOfRangeSnafu, DataFusionPlanningSnafu, Result, TimeIndexNotFoundSnafu, +}; +use crate::promql::planner::PromPlanner; +use crate::query_engine::QueryEngineState; + +impl PromPlanner { + /// Resolve the `@` modifier of a vector or matrix selector into the timestamp its sample + /// window is anchored at, in milliseconds since the Unix epoch. + /// + /// Prometheus semantics: + /// - `@ ` anchors at the given timestamp, + /// - `@ start()` / `@ end()` anchor at the evaluation range of the whole statement, + /// - `offset` shifts the anchor backwards: the window ends at `anchor - offset`. + /// + /// Returns `None` when the selector has no `@` modifier. + pub(crate) fn at_ref_time( + &self, + at: &Option, + offset: &Option, + ) -> Result> { + let anchor = match at { + None => return Ok(None), + Some(AtModifier::Start) => self.ctx.stmt_start, + Some(AtModifier::End) => self.ctx.stmt_end, + Some(AtModifier::At(time)) => Self::system_time_to_millis(time)?, + }; + Ok(Some(Self::anchor_sub(anchor, Self::offset_millis(offset))?)) + } + + /// Subtracts `rhs` from `lhs` on the millisecond timeline of an `@` anchor. + /// + /// A negative result is valid: `@` and `offset` accept timestamps before the Unix epoch. A + /// result outside the representable millisecond range is rejected like an unrepresentable + /// anchor ([`Self::system_time_to_millis`]) instead of clamping it, so the same class of + /// input always gets the same answer. + pub(crate) fn anchor_sub(lhs: Millisecond, rhs: Millisecond) -> Result { + lhs.checked_sub(rhs) + .with_context(|| AtModifierTimestampOutOfRangeSnafu { + timestamp: format!("{}ms - {}ms", lhs, rhs), + }) + } + + /// The offset a selector with an `@` modifier is evaluated with. + /// + /// Prometheus anchors such a selector by rewriting its offset to `eval_time - anchor` + /// (`setOffsetForAtModifier`), so that the selector always selects its samples around `anchor` + /// regardless of the step being evaluated. `eval_time` is the start of the evaluation the + /// selector belongs to, which is `ctx.start`. + /// + /// Returns `None` when the selector has no `@` modifier. + pub(crate) fn at_modifier_offset( + &self, + at: &Option, + offset: &Option, + ) -> Result> { + let Some(anchor) = self.at_ref_time(at, offset)? else { + return Ok(None); + }; + Ok(Some(Self::anchor_sub(self.ctx.start, anchor)?)) + } + + /// Whether `expr` is a call that has to be evaluated once for the whole grid, because it folds + /// a range selector anchored by `@` — the range argument of the call's parser signature, which + /// is a [`MatrixSelector`] here; only a call with such an argument can take this path, so no + /// function-name registry is involved. + /// + /// This is the shape that needs Prometheus' `StepInvariantExpr` the most: a range function such + /// as `rate` derives its result from the evaluation instant it is called at, so folding the + /// window once per step would let the outer evaluation grid change the result of a window that + /// `@` fixed. Evaluated once, at the start of the grid, the rewritten offset of + /// [`Self::at_modifier_offset`] places the anchor at that instant, and [`Self::replay_over_grid`] + /// reports the result at every step. + /// + /// Unlike Prometheus, which wraps the whole step-invariant subtree (`preprocessExprHelper`), + /// only the call itself is promoted here. The operators above it are not: they are still planned + /// at every step over the replayed result, which is safe for the row-wise ones and keeps the + /// promotion root narrow. The promotion root has to stay a call over one range selector, because + /// the replay needs one series per batch ([`Self::series_divide_plan`]) and only such a call + /// guarantees that the rows it emits still describe the series it was divided by. The operators + /// left out — an aggregation, a join, or a label rewriting call such as `label_join` — mix or + /// re-label the rows of different series, so they are unsafe as promotion roots even though + /// evaluating them after the replay is fine. A call whose input is an anchored *instant* + /// selector (`abs(some_metric @ 300)`) needs no promotion either: the selector anchors and + /// replays its sample per series on its own, and the call above it is row-wise. + /// + /// `predict_linear` is the exception among the range functions: it predicts from the evaluation + /// instant of each step ([`Self::create_range_eval_ts_expr`]), so it has to stay outside the + /// promoted subtree and follow the grid. The remaining arguments of the call have to be + /// literals, since the replay of the promoted result has no second vector input to divide. + /// Parentheses around the range argument are transparent (`rate((m[5m] @ 300))`), so they are + /// looked through and the call is promoted as if they were absent. Nothing else of the subtree + /// is unwrapped, so an outer parenthesis promotes no operator above the call. + fn promotes_anchored_range_call(expr: &PromExpr) -> bool { + let PromExpr::Call(Call { func, args }) = expr else { + return false; + }; + // See the doc comment: the regression of `predict_linear` follows the evaluation step. + if func.name == "predict_linear" { + return false; + } + let mut anchored_range = false; + for arg in &args.args { + // Parentheses around the range argument are transparent, so the call is promoted the + // same way for `rate((m[5m] @ 300))` as for `rate(m[5m] @ 300)`. Only the parentheses + // directly around this one argument are looked through here: the promotion stays + // confined to a call over one anchored range selector instead of descending into an + // arbitrary parenthesized subtree. + let mut arg = arg.as_ref(); + while let PromExpr::Paren(ParenExpr { expr }) = arg { + arg = expr; + } + match arg { + // The window is pinned by `@`, so every step folds the same samples. + PromExpr::MatrixSelector(MatrixSelector { vs, .. }) if vs.at.is_some() => { + if anchored_range { + return false; + } + anchored_range = true; + } + // A literal argument is the same value at every step. + arg if Self::try_build_literal_expr(arg).is_some() => {} + _ => return false, + } + } + anchored_range + } + + /// Plans the anchored range call `prom_expr` ([`Self::promotes_anchored_range_call`]) as a + /// step-invariant subtree: the call is evaluated on a single evaluation instant (`grid_start`, + /// the start of the outer evaluation) and its result is then reported at every step of + /// `[grid_start, ctx.end]` by [`Self::replay_over_grid`]. + /// + /// This is the planner's counterpart of Prometheus' `StepInvariantExpr` for the one shape it + /// promotes. Evaluating the call once matters for the functions that derive their result from + /// the step being evaluated: `rate(m[5m] @ 300)` folds its window around the anchor once, and + /// the extrapolation boundaries of `rate` must be derived from that same window at every step + /// instead of following the outer evaluation timestamp. + /// + /// Only the call itself is promoted; the operators above it are planned as usual over the + /// replayed result. The result of the promoted call is split into one series per batch before it + /// is replayed ([`Self::series_divide_plan`]), because the row-wise projection of the call does + /// not preserve the batch layout of the selector. + /// + /// Returns `None` when `prom_expr` is not such a call, so that the caller plans it as usual. + /// The selector inside the promoted call keeps its own `@` anchoring (see + /// [`Self::at_modifier_offset`]), and planning it with `ctx.end == ctx.start` folds its window + /// once for that single instant instead of expanding it over the grid, which the replay of the + /// call result above already does. + pub(crate) async fn promote_anchored_range_call( + &mut self, + prom_expr: &PromExpr, + timestamp_fn: bool, + query_engine_state: &QueryEngineState, + ) -> Result> { + let grid_start = self.ctx.start; + let grid_end = self.ctx.end; + // An instant query evaluates a single step, so there is nothing to promote. + if grid_start == grid_end || !Self::promotes_anchored_range_call(prom_expr) { + return Ok(None); + } + + // Plan the subtree on a single evaluation instant: every selector below still anchors its + // window through `@`, and the functions above them derive their result from that one + // instant. The planner is single-use, so `ctx.end` needs no restore-on-error. + self.ctx.end = grid_start; + let anchored = self + .prom_expr_to_plan_inner(prom_expr, timestamp_fn, query_engine_state) + .await?; + self.ctx.end = grid_end; + + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + // The replay reads one series per batch; see [`Self::series_divide_plan`] for why the + // layout of the selector below the promoted call does not survive it. + let anchored = self.series_divide_plan(anchored, &time_index_column)?; + Ok(Some(self.replay_over_grid( + anchored, + grid_start, + grid_end, + time_index_column, + ))) + } + + /// Convert the timestamp of an `@` modifier into milliseconds since the Unix epoch. + fn system_time_to_millis(time: &SystemTime) -> Result { + let (millis, negative) = match time.duration_since(UNIX_EPOCH) { + Ok(duration) => (duration.as_millis(), false), + // The `@` modifier accepts timestamps before the Unix epoch, e.g. `@ -1`. + Err(err) => (err.duration().as_millis(), true), + }; + ensure!( + millis <= i64::MAX as u128, + AtModifierTimestampOutOfRangeSnafu { + timestamp: time + .duration_since(UNIX_EPOCH) + .map(|duration| format!("+{}ms", duration.as_millis())) + .unwrap_or_else(|err| format!("-{}ms", err.duration().as_millis())), + } + ); + let millis = millis as Millisecond; + Ok(if negative { -millis } else { millis }) + } + + /// Report the samples of `anchored` at every step of the evaluation grid + /// `[grid_start, grid_end]`. + /// + /// A selector with an `@` modifier is anchored: the sample window is selected once, around the + /// anchor timestamp, and every evaluation step reports that same window. Prometheus does this + /// by rewriting the selector's offset to `eval_time - anchor` and only fetching the samples on + /// the first step (`setOffsetForAtModifier` plus the `refetch` shortcut in `rangeEval`). + /// + /// The expansion reuses [`InstantManipulate`] with a lookback that spans the whole grid: every + /// step then picks the same sample (or the same already computed value, when `anchored` ends + /// with a function call such as `rate`) and stamps it with the step's timestamp. + /// + /// Every input batch of `anchored` must hold exactly one series, because [`InstantManipulate`] + /// takes a batch as one timeline. A leaf-level replay (`m @ 300`) consumes the [`SeriesDivide`] + /// of its selector directly. A promoted call is guaranteed that layout by + /// [`Self::series_divide_plan`], which is why [`Self::promote_anchored_range_call`] splits + /// its result before calling this method. + pub(crate) fn replay_over_grid( + &self, + anchored: LogicalPlan, + grid_start: Millisecond, + grid_end: Millisecond, + time_index_column: String, + ) -> LogicalPlan { + if grid_start == grid_end { + // A single evaluation step: `anchored` is already stamped with that timestamp. + return anchored; + } + + let series_key_columns = self.series_key_columns_for_schema(anchored.schema()); + let replayed = InstantManipulate::new( + grid_start, + grid_end, + // The lookback must keep the single anchored sample eligible for every step. + grid_end - grid_start + 1, + self.ctx.interval, + 0, + time_index_column, + series_key_columns, + self.ctx.field_columns.first().cloned(), + anchored, + ); + LogicalPlan::Extension(Extension { + node: Arc::new(replayed), + }) + } + + /// Sorts `input` by its series key and time index and splits it into one batch per series. + /// + /// [`InstantManipulate`] reads every input batch as one series (it takes the timeline of the + /// batch and reports the row selected at every step), so a batch holding several series would + /// lose all but one of them. A selector establishes that layout with its own [`SeriesDivide`], + /// but the per-series distribution requirement does not reach a promoted subtree above it: + /// the row-wise projection of a call sits in between, so the batch boundaries of the selector + /// are not preserved — in a distributed plan the promoted result can be delivered as one batch + /// holding every series. Sorting and dividing here restores the layout, exactly like + /// [`Self::prom_matrix_selector_to_plan`] does for the input of a range function. + /// + /// Series keys that are not present in `input` are dropped, since `ctx.tag_columns` may have + /// drifted from the actual output schema. A plan without any series key column is returned + /// unchanged: there is nothing to divide by. + fn series_divide_plan( + &self, + input: LogicalPlan, + time_index_column: &str, + ) -> Result { + let series_key_columns = self.series_key_columns_for_schema(input.schema()); + if series_key_columns.is_empty() { + return Ok(input); + } + + let mut sort_exprs = series_key_columns + .iter() + .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true)) + .collect::>(); + sort_exprs.push(DfExpr::Column(Column::from_name(time_index_column)).sort(true, true)); + let sort_plan = LogicalPlanBuilder::from(input) + .sort(sort_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + series_key_columns, + time_index_column.to_string(), + sort_plan, + )), + })) + } +} diff --git a/src/query/src/promql/planner/function_plans.rs b/src/query/src/promql/planner/function_plans.rs new file mode 100644 index 00000000000..4723fe35232 --- /dev/null +++ b/src/query/src/promql/planner/function_plans.rs @@ -0,0 +1,473 @@ +// 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. + +//! Dedicated helper plans for PromQL functions that cannot be expressed as a +//! plain scalar function over the input plan: the histogram helpers, `vector()`, +//! `scalar()`, and `absent()`. + +use std::sync::Arc; + +use common_query::prelude::greptime_value; +use datafusion::functions_aggregate::expr_fn::first_value; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; +use datafusion::prelude::{Column, Expr as DfExpr}; +use datafusion::scalar::ScalarValue; +use datafusion_common::DFSchema; +use datafusion_expr::expr_fn::when; +use datafusion_expr::simplify::SimplifyContext; +use datafusion_expr::{ident, lit}; +use datafusion_functions::core::coalesce; +use datatypes::arrow::datatypes::DataType as ArrowDataType; +use promql::extension_plan::{ + Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, ScalarCalculate, +}; +use promql::functions::{NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile}; +use promql_parser::label::MatchOp; +use promql_parser::parser::{Expr as PromExpr, FunctionArgs as PromFunctionArgs}; +use snafu::{OptionExt, ResultExt, ensure}; + +use crate::promql::error::{ + DataFusionPlanningSnafu, FunctionInvalidArgumentSnafu, MultiFieldsNotSupportedSnafu, + PromqlPlanNodeSnafu, Result, TimeIndexNotFoundSnafu, ValueNotFoundSnafu, +}; +use crate::promql::planner::{ + LE_COLUMN_NAME, PromPlanner, SCALAR_FUNCTION, SPECIAL_ABSENT_FUNCTION, + SPECIAL_HISTOGRAM_FRACTION, SPECIAL_HISTOGRAM_QUANTILE, SPECIAL_TIME_FUNCTION, + SPECIAL_VECTOR_FUNCTION, +}; +use crate::query_engine::QueryEngineState; + +impl PromPlanner { + /// Create a classic, native, or mixed histogram helper plan. + pub(super) async fn create_histogram_plan( + &mut self, + function_name: &str, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + let float_literal = |param: &PromExpr| -> Result { + let value = (|| { + let expr = Self::get_param_as_literal_expr( + Some(param), + None, + Some(ArrowDataType::Float64), + ) + .ok()?; + let simplifier = ExprSimplifier::new(SimplifyContext::default()); + let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; + let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { + return None; + }; + let ScalarValue::Float64(Some(value)) = + value.cast_to(&ArrowDataType::Float64).ok()? + else { + return None; + }; + Some(value) + })() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + })?; + Ok(value) + }; + let (function, input) = match (function_name, args.args.as_slice()) { + (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( + HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), + input.as_ref().clone(), + ), + (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( + HistogramFoldOperation::Fraction { + lower: float_literal(lower)?.into(), + upper: float_literal(upper)?.into(), + }, + input.as_ref().clone(), + ), + _ => { + return FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + } + .fail(); + } + }; + + let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; + // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a + // stable series identifier anymore. HistogramFold must not treat it as a label column. + let input_plan = self.strip_tsid_column(input_plan)?; + self.ctx.use_tsid = false; + + if let Some((float_field, histogram_field)) = + Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) + { + if self.ctx.has_le_tag() { + return self.create_mixed_histogram_plan( + function, + input_plan, + float_field, + histogram_field, + ); + } + self.ctx.field_columns = vec![histogram_field]; + } + if self.all_field_columns_are_native_histograms(input_plan.schema()) { + return self.create_native_histogram_plan(function, input_plan); + } + + if !self.ctx.has_le_tag() { + // Return empty result instead of error when 'le' column is not found + // This handles the case when histogram metrics don't exist + return Ok(LogicalPlan::EmptyRelation( + datafusion::logical_expr::EmptyRelation { + produce_one_row: false, + schema: input_plan.schema().clone(), + }, + )); + } + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + // FIXME(ruihang): support multi fields + let field_column = self + .ctx + .field_columns + .first() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function.function_name().to_string(), + })? + .clone(); + // remove le column from tag columns + self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); + + let fold = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + field_column, + time_index_column, + function, + None, + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(fold), + })) + } + + fn create_native_histogram_expr( + &self, + function: HistogramFoldOperation, + field_column: &str, + ) -> DfExpr { + let field = DfExpr::Column(Column::from_name(field_column)); + let (func, args) = match function { + HistogramFoldOperation::Quantile(quantile) => ( + Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(quantile))], + ), + HistogramFoldOperation::Fraction { lower, upper } => ( + Arc::new(NativeHistogramFraction::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(lower)), lit(f64::from(upper))], + ), + }; + DfExpr::ScalarFunction(ScalarFunction { func, args }) + } + + fn create_native_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + ) -> Result { + ensure!( + self.ctx.field_columns.len() == 1, + MultiFieldsNotSupportedSnafu { + operator: function.function_name() + }, + ); + + let field_column = self.ctx.field_columns[0].clone(); + let function_expr = self.create_native_histogram_expr(function, &field_column); + let display_name = function_expr.schema_name().to_string(); + self.ctx.field_columns = vec![display_name.clone()]; + + let project_exprs = std::iter::once(self.create_time_index_column_expr()?) + .chain(std::iter::once(function_expr.alias(display_name))) + .chain(self.create_tag_column_exprs()?) + .collect::>(); + + LogicalPlanBuilder::from(input_plan) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .filter(self.create_empty_values_filter_expr(false)?) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + fn create_mixed_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + float_field: String, + histogram_field: String, + ) -> Result { + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let tag_columns = self.ctx.tag_columns.clone(); + let folded = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + float_field.clone(), + time_index_column.clone(), + function, + Some(histogram_field.clone()), + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + let record_collision = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( + "vector contains a mix of classic and native histograms".to_string(), + self.promql_annotations.clone(), + )), + args: vec![ident(&float_field), ident(&histogram_field)], + }); + let keep = when( + ident(&float_field) + .is_not_null() + .and(ident(&histogram_field).is_not_null()), + record_collision, + ) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu)?; + + let native_expr = self.create_native_histogram_expr(function, &histogram_field); + let output_field = native_expr.schema_name().to_string(); + let value = DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![ident(&float_field), native_expr], + }); + self.ctx.field_columns = vec![output_field.clone()]; + LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { + node: Arc::new(folded), + })) + .filter(keep) + .context(DataFusionPlanningSnafu)? + .project( + std::iter::once(ident(&time_index_column)) + .chain(std::iter::once(value.alias(output_field))) + .chain(tag_columns.iter().map(ident)), + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + /// Create a [SPECIAL_VECTOR_FUNCTION] plan + pub(super) async fn create_vector_plan( + &mut self, + args: &PromFunctionArgs, + ) -> Result { + if args.args.len() != 1 { + return FunctionInvalidArgumentSnafu { + fn_name: SPECIAL_VECTOR_FUNCTION.to_string(), + } + .fail(); + } + let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; + + // reuse `SPECIAL_TIME_FUNCTION` as name of time index column + self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); + self.ctx.reset_table_name_and_schema(); + self.ctx.tag_columns = vec![]; + self.ctx.aggregation_field_labels.clear(); + self.ctx.field_columns = vec![greptime_value().to_string()]; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new( + EmptyMetric::new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + SPECIAL_TIME_FUNCTION.to_string(), + greptime_value().to_string(), + Some(lit), + ) + .context(DataFusionPlanningSnafu)?, + ), + })) + } + + /// Create a [SCALAR_FUNCTION] plan + pub(super) async fn create_scalar_plan( + &mut self, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + ensure!( + args.len() == 1, + FunctionInvalidArgumentSnafu { + fn_name: SCALAR_FUNCTION + } + ); + let input = self + .prom_expr_to_plan(&args.args[0], query_engine_state) + .await?; + let input_schema = input.schema().clone(); + let alternative_samples = + Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns); + let histogram_fields = self + .ctx + .field_columns + .iter() + .filter(|field| Self::field_column_is_native_histogram(&input_schema, field)) + .count(); + ensure!( + self.ctx.field_columns.len() == 1 || alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: SCALAR_FUNCTION + }, + ); + let scalar_field = self + .ctx + .field_columns + .iter() + .find(|field| !Self::field_column_is_native_histogram(&input_schema, field)) + .or_else(|| self.ctx.field_columns.first()) + .cloned() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: SCALAR_FUNCTION, + })?; + let input = if histogram_fields == self.ctx.field_columns.len() { + // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN + // for every evaluation timestamp without attempting a Struct-to-Float64 cast. + LogicalPlanBuilder::from(input) + .filter(lit(false)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else if histogram_fields > 0 { + // A mixed vector contributes only its float samples to scalar(). + LogicalPlanBuilder::from(input) + .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + input + }; + let scalar_plan = LogicalPlan::Extension(Extension { + node: Arc::new( + ScalarCalculate::new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + input, + self.ctx.time_index_column.as_ref().unwrap(), + &self.ctx.tag_columns, + &scalar_field, + self.ctx.table_name.as_deref(), + ) + .context(PromqlPlanNodeSnafu)?, + ), + }); + // scalar plan have no tag columns + self.ctx.tag_columns.clear(); + self.ctx.aggregation_field_labels.clear(); + self.ctx.field_columns.clear(); + self.ctx + .field_columns + .push(scalar_plan.schema().field(1).name().clone()); + Ok(scalar_plan) + } + + /// Create a [SPECIAL_ABSENT_FUNCTION] plan + pub(super) async fn create_absent_plan( + &mut self, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + if args.args.len() != 1 { + return FunctionInvalidArgumentSnafu { + fn_name: SPECIAL_ABSENT_FUNCTION.to_string(), + } + .fail(); + } + let input = self + .prom_expr_to_plan(&args.args[0], query_engine_state) + .await?; + + let time_index_expr = self.create_time_index_column_expr()?; + let first_field_expr = + self.create_field_column_exprs()? + .pop() + .with_context(|| ValueNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let first_value_expr = first_value(first_field_expr, vec![]); + + let ordered_aggregated_input = LogicalPlanBuilder::from(input) + .aggregate( + vec![time_index_expr.clone()], + vec![first_value_expr.clone()], + ) + .context(DataFusionPlanningSnafu)? + .sort(vec![time_index_expr.sort(true, false)]) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + let fake_labels = self + .ctx + .selector_matcher + .iter() + .filter_map(|matcher| match matcher.op { + MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())), + _ => None, + }) + .collect::>(); + + // Create the absent plan + let absent_plan = LogicalPlan::Extension(Extension { + node: Arc::new( + Absent::try_new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + self.ctx.time_index_column.as_ref().unwrap().clone(), + self.ctx.field_columns[0].clone(), + fake_labels, + ordered_aggregated_input, + ) + .context(DataFusionPlanningSnafu)?, + ), + }); + + // The absent series carries the equality matchers as labels, not the input's + // tags or value fields, so the input's field grouping labels no longer apply. + self.ctx.aggregation_field_labels.clear(); + Ok(absent_plan) + } +} diff --git a/src/query/src/promql/planner/island.rs b/src/query/src/promql/planner/island.rs new file mode 100644 index 00000000000..d7c04f75015 --- /dev/null +++ b/src/query/src/promql/planner/island.rs @@ -0,0 +1,536 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The binary-island fast path of the PromQL planner. + +use std::collections::{BTreeSet, HashMap}; + +use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL; +use datafusion::common::DFSchemaRef; +use datafusion::logical_expr::expr::Alias; +use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; +use datafusion_common::{NullEquality, TableReference}; +use datafusion_expr::lit; +use promql_parser::label::{METRIC_NAME, MatchOp, Matcher}; +use promql_parser::parser::token::{self, TokenType}; +use promql_parser::parser::{ + BinaryExpr as PromBinaryExpr, Expr as PromExpr, Offset, ParenExpr, UnaryExpr, + VectorMatchCardinality, VectorSelector, +}; +use snafu::ResultExt; + +use crate::promql::error::{DataFusionPlanningSnafu, Result}; +use crate::promql::planner::{PromPlanner, PromPlannerContext}; + +/// Prefix for generated binary island leaf aliases. +const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct VectorLeafKey { + metric_name: String, + matchers: Vec<(String, String, String)>, + or_matchers: Vec>, + offset_ms: i128, + at: String, +} + +#[derive(Debug, Clone)] +struct IslandLeaf { + selector: VectorSelector, + display_table: String, +} + +#[derive(Debug, Clone)] +enum IslandExpr { + VectorLeaf(usize), + Scalar(DfExpr), + Unary { + input: Box, + }, + Binary { + op: TokenType, + lhs: Box, + rhs: Box, + }, +} + +impl IslandExpr { + fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option { + if let Some(expr) = PromPlanner::try_build_literal_expr(expr) { + return Some(Self::Scalar(expr)); + } + + match expr { + PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env), + PromExpr::VectorSelector(selector) => { + let leaf = env.intern_leaf(selector)?; + Some(Self::VectorLeaf(leaf)) + } + PromExpr::Unary(UnaryExpr { expr }) => { + let input = Self::try_new(expr, env)?; + Some(Self::Unary { + input: Box::new(input), + }) + } + PromExpr::Binary(PromBinaryExpr { + lhs, + rhs, + op, + modifier, + }) if matches!( + op.id(), + token::T_ADD + | token::T_SUB + | token::T_MUL + | token::T_DIV + | token::T_MOD + | token::T_POW + | token::T_ATAN2 + ) && modifier.as_ref().is_none_or(|modifier| { + !modifier.return_bool + && modifier.matching.is_none() + && matches!(modifier.card, VectorMatchCardinality::OneToOne) + && modifier.fill_values.lhs.is_none() + && modifier.fill_values.rhs.is_none() + }) => + { + let lhs = Self::try_new(lhs, env)?; + let rhs = Self::try_new(rhs, env)?; + Some(Self::Binary { + op: *op, + lhs: Box::new(lhs), + rhs: Box::new(rhs), + }) + } + _ => None, + } + } +} + +#[derive(Debug, Default)] +struct IslandCollectEnv { + leaf_by_key: HashMap, + leaves: Vec, + vector_occurrences: usize, +} + +#[derive(Debug)] +struct PlannedIslandLeaf { + plan: LogicalPlan, + ctx: PromPlannerContext, + alias: TableReference, + display_table: String, +} + +#[derive(Debug)] +struct IslandFieldExprs { + exprs: Vec, + names: Vec, + scalar: bool, +} + +impl VectorLeafKey { + fn from_selector(selector: &VectorSelector) -> Option { + let mut metric_name = selector.name.clone(); + let mut matchers = Vec::with_capacity(selector.matchers.matchers.len()); + let matcher_key = |matcher: &Matcher| { + ( + matcher.name.clone(), + matcher.op.to_string(), + matcher.value.clone(), + ) + }; + + for matcher in &selector.matchers.matchers { + if matcher.name == METRIC_NAME { + if matcher.op != MatchOp::Equal || metric_name.is_some() { + return None; + } + metric_name = Some(matcher.value.clone()); + } else { + matchers.push(matcher_key(matcher)); + } + } + matchers.sort(); + + let mut or_matchers = selector + .matchers + .or_matchers + .iter() + .map(|group| { + let mut group = group.iter().map(matcher_key).collect::>(); + group.sort(); + group + }) + .collect::>(); + or_matchers.sort(); + + Some(Self { + metric_name: metric_name?, + matchers, + or_matchers, + offset_ms: match &selector.offset { + Some(Offset::Pos(duration)) => duration.as_millis() as i128, + Some(Offset::Neg(duration)) => -(duration.as_millis() as i128), + None => 0, + }, + at: format!("{:?}", selector.at), + }) + } +} + +impl IslandCollectEnv { + fn intern_leaf(&mut self, selector: &VectorSelector) -> Option { + self.vector_occurrences += 1; + let key = VectorLeafKey::from_selector(selector)?; + if let Some(id) = self.leaf_by_key.get(&key) { + return Some(*id); + } + + let id = self.leaves.len(); + self.leaves.push(IslandLeaf { + selector: selector.clone(), + display_table: key.metric_name.clone(), + }); + self.leaf_by_key.insert(key, id); + Some(id) + } +} + +impl PromPlanner { + pub(super) async fn try_plan_binary_island( + &mut self, + binary_expr: &PromBinaryExpr, + ) -> Result> { + let original_ctx = self.ctx.clone(); + let mut collect_env = IslandCollectEnv::default(); + let Some(island_expr) = + IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env) + else { + return Ok(None); + }; + + if collect_env.leaves.is_empty() + || collect_env.vector_occurrences <= collect_env.leaves.len() + { + return Ok(None); + } + + let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len()); + for (idx, leaf) in collect_env.leaves.iter().enumerate() { + let plan = self + .prom_vector_selector_to_plan(&leaf.selector, false) + .await?; + let ctx = self.ctx.clone(); + let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}")); + let plan = LogicalPlanBuilder::from(plan) + .alias(alias.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + planned_leaves.push(PlannedIslandLeaf { + plan, + ctx, + alias, + display_table: leaf.display_table.clone(), + }); + } + + if planned_leaves.iter().any(|leaf| { + Self::field_columns_contain_native_histogram( + leaf.plan.schema(), + &leaf.ctx.field_columns, + ) + }) { + self.ctx = original_ctx; + return Ok(None); + } + + if !Self::binary_island_join_contexts_supported(&planned_leaves) { + self.ctx = original_ctx; + return Ok(None); + } + + let mut input = planned_leaves[0].plan.clone(); + for right_idx in 1..planned_leaves.len() { + input = self.join_binary_island_leaf( + input, + &planned_leaves[0], + &planned_leaves[right_idx], + )?; + } + + let field_exprs = + Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?; + if field_exprs.scalar || field_exprs.exprs.is_empty() { + self.ctx = original_ctx; + return Ok(None); + } + + let plan = self.project_binary_island( + input, + &planned_leaves[0].alias, + &planned_leaves[0].ctx, + field_exprs, + )?; + Ok(Some(plan)) + } + + fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool { + if leaves + .iter() + .any(|leaf| leaf.ctx.time_index_column.is_none()) + { + return false; + } + + if leaves.len() <= 1 { + return true; + } + + let first_tags = leaves[0].ctx.tag_columns.iter().collect::>(); + + leaves.iter().skip(1).all(|leaf| { + (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan)) + || leaf.ctx.tag_columns.iter().collect::>() == first_tags + }) + } + + fn join_binary_island_leaf( + &self, + left: LogicalPlan, + first_leaf: &PlannedIslandLeaf, + right_leaf: &PlannedIslandLeaf, + ) -> Result { + let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty() + || right_leaf.ctx.tag_columns.is_empty()) + && !first_leaf + .ctx + .tag_columns + .iter() + .chain(&right_leaf.ctx.tag_columns) + .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL); + let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns( + left.schema(), + right_leaf.plan.schema(), + &first_leaf.ctx, + &right_leaf.ctx, + only_join_time_index, + &None, + )?; + + if let (Some(left_time_index_column), Some(right_time_index_column)) = ( + first_leaf.ctx.time_index_column.clone(), + right_leaf.ctx.time_index_column.clone(), + ) { + left_keys.insert(left_time_index_column); + right_keys.insert(right_time_index_column); + } + + LogicalPlanBuilder::from(left) + .join_detailed( + right_leaf.plan.clone(), + JoinType::Inner, + ( + left_keys + .into_iter() + .map(|name| Column::new(Some(first_leaf.alias.clone()), name)) + .collect::>(), + right_keys + .into_iter() + .map(|name| Column::new(Some(right_leaf.alias.clone()), name)) + .collect::>(), + ), + force_empty_join.then_some(lit(false)), + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + fn build_binary_island_field_exprs( + expr: &IslandExpr, + leaves: &[PlannedIslandLeaf], + schema: &DFSchemaRef, + ) -> Result { + match expr { + IslandExpr::VectorLeaf(id) => { + let leaf = &leaves[*id]; + let exprs = leaf + .ctx + .field_columns + .iter() + .map(|field| { + schema + .qualified_field_with_name(Some(&leaf.alias), field) + .context(DataFusionPlanningSnafu) + .map(|field| DfExpr::Column(field.into())) + }) + .collect::>>()?; + let names = leaf + .ctx + .field_columns + .iter() + .map(|field| format!("{}.{}", leaf.display_table, field)) + .collect(); + Ok(IslandFieldExprs { + exprs, + names, + scalar: false, + }) + } + IslandExpr::Scalar(expr) => Ok(IslandFieldExprs { + exprs: vec![expr.clone()], + names: vec![expr.schema_name().to_string()], + scalar: true, + }), + IslandExpr::Unary { input } => { + let input = Self::build_binary_island_field_exprs(input, leaves, schema)?; + let mut exprs = Vec::with_capacity(input.exprs.len()); + let mut names = Vec::with_capacity(input.names.len()); + for (expr, name) in input.exprs.into_iter().zip(input.names) { + exprs.push(DfExpr::Negative(Box::new(expr))); + names.push(format!("-{name}")); + } + Ok(IslandFieldExprs { + exprs, + names, + scalar: input.scalar, + }) + } + IslandExpr::Binary { op, lhs, rhs } => { + let same_leaf = match (&**lhs, &**rhs) { + (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right)) + if left == right => + { + Some(*left) + } + _ => None, + }; + let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?; + let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?; + let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let scalar = lhs.scalar && rhs.scalar; + let op = op.to_string(); + + let (exprs, names) = match (lhs.scalar, rhs.scalar) { + (true, true) => { + let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?; + let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]); + (vec![expr], vec![name]) + } + (true, false) => { + let mut exprs = Vec::with_capacity(rhs.exprs.len()); + let mut names = Vec::with_capacity(rhs.names.len()); + for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) { + exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?); + names.push(format!("{} {op} {rhs_name}", lhs.names[0])); + } + (exprs, names) + } + (false, true) => { + let mut exprs = Vec::with_capacity(lhs.exprs.len()); + let mut names = Vec::with_capacity(lhs.names.len()); + for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) { + exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?); + names.push(format!("{lhs_name} {op} {}", rhs.names[0])); + } + (exprs, names) + } + (false, false) => { + let mut exprs = Vec::new(); + let mut names = Vec::new(); + for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs + .exprs + .into_iter() + .zip(rhs.exprs) + .zip(lhs.names.into_iter().zip(rhs.names)) + .enumerate() + { + if let Some(leaf) = same_leaf { + let field = leaves[leaf] + .ctx + .field_columns + .get(idx) + .cloned() + .unwrap_or_else(|| lhs_name.clone()); + lhs_name = format!("lhs.{field}"); + rhs_name = format!("rhs.{field}"); + } + exprs.push(expr_builder(lhs_expr, rhs_expr)?); + names.push(format!("{lhs_name} {op} {rhs_name}")); + } + (exprs, names) + } + }; + + Ok(IslandFieldExprs { + exprs, + names, + scalar, + }) + } + } + } + + fn project_binary_island( + &mut self, + input: LogicalPlan, + base_alias: &TableReference, + base_ctx: &PromPlannerContext, + field_exprs: IslandFieldExprs, + ) -> Result { + self.ctx = base_ctx.clone(); + + let schema = input.schema(); + let non_field_exprs = base_ctx + .tag_columns + .iter() + .chain(base_ctx.time_index_column.iter()) + .map(|column| { + schema + .qualified_field_with_name(Some(base_alias), column) + .context(DataFusionPlanningSnafu) + .map(|field| DfExpr::Column(field.into())) + }); + let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid) + .into_iter() + .map(Ok); + + self.ctx.field_columns = field_exprs.names; + let field_exprs = field_exprs + .exprs + .into_iter() + .zip(self.ctx.field_columns.iter()) + .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::, name)))); + + let project_exprs = non_field_exprs + .chain(tsid_expr) + .chain(field_exprs) + .collect::>>()?; + + let plan = LogicalPlanBuilder::from(input) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + self.ctx.table_name = None; + self.ctx.schema_name = None; + + Ok(plan) + } +} diff --git a/src/query/src/promql/planner/matching_filters.rs b/src/query/src/promql/planner/matching_filters.rs index 386d95e9228..d70e1bc4192 100644 --- a/src/query/src/promql/planner/matching_filters.rs +++ b/src/query/src/promql/planner/matching_filters.rs @@ -52,26 +52,24 @@ const LABEL_PRESERVING_RANGE_FUNCTIONS: [&str; 20] = [ /// Copies label matchers from one operand of `binary` to the other, so both scans discard /// non-joining series instead of the join. /// -/// One-to-one arithmetic operands are inner-joined on their matching labels with plain column +/// Arithmetic operands are inner-joined on their matching labels with plain column /// equality, so a matcher on a matching label already holds for every surviving pair. Copying it /// to the other operand can only drop rows that had no partner, whatever the matcher kind, and /// `NULL` pairs only with `NULL`. Nor can it split a match group, whose series agree on every -/// matching label, so it does not affect a cardinality check such as #9209. +/// matching label: a propagated matcher removes whole groups, so the cardinality check still +/// sees every group that takes part in the matching. A duplicate in a group without a partner +/// goes unreported as a result, where Prometheus fails the query on it. +/// Grouped matching is conservatively restricted to a provably unique one-side operand. /// /// `left_tags` and `right_tags` are the tag columns of the planned operands: a matcher may just /// as well constrain a value field, which the parsed expression does not distinguish from a -/// label. For an aggregated operand they are its grouping labels, which is what makes filtering -/// its input equivalent to filtering its output. +/// label. Matchers can cross an aggregate only on labels that partition its input. +/// `left_field_labels` and `right_field_labels` exclude value fields promoted to grouping +/// labels: filtering them before sample selection could replace the latest sample with an +/// older one (#9242). /// -/// A value field is not a series identity though: its value may vary between the samples of one -/// series, so a matcher on it must not be lowered below sample selection (`PromInstantManipulate`, -/// `PromSeriesDivide`), where discarding a sample promotes an older one out of the lookback window -/// and fabricates a match. An aggregated operand's tag columns are its grouping labels rather than -/// the tags of its input, so those labels may name value fields; a matcher therefore only crosses -/// an aggregation when it names one of the grouping labels that aggregation partitions by, and -/// `left_field_labels`/`right_field_labels` rule out the grouping labels the planner found to be -/// value fields of the aggregated operand's input -/// (`PromPlannerContext::aggregation_field_labels`). +/// Whether a matcher reaches a scan is decided per receiving operand ([`preserves_filter`]), so +/// propagation is one-way when only one of the two can take the filter. pub(super) fn propagate( binary: &BinaryExpr, left_tags: &[String], @@ -108,35 +106,16 @@ fn try_propagate( ) { return Err("operator is not arithmetic"); } - if let Some(modifier) = &binary.modifier { - if !matches!(modifier.card, VectorMatchCardinality::OneToOne) { - return Err("matching is not one-to-one"); - } - if modifier.fill_values.lhs.is_some() || modifier.fill_values.rhs.is_some() { - return Err("operand carries a fill modifier"); - } + if let Some(modifier) = &binary.modifier + && (modifier.fill_values.lhs.is_some() || modifier.fill_values.rhs.is_some()) + { + return Err("operand carries a fill modifier"); } let matching = binary .modifier .as_ref() .and_then(|modifier| modifier.matching.as_ref()); - // The grouping labels of every aggregation the operands cross, collected before `rewritten` - // takes the mutable borrows that rule out borrowing `binary` again. - let left_grouped = modifier_label_names(&binary.lhs); - let left_ignored = modifier_excluded_label_names(&binary.lhs); - let right_grouped = modifier_label_names(&binary.rhs); - let right_ignored = modifier_excluded_label_names(&binary.rhs); - - let mut rewritten = binary.clone(); - let (left, left_restricted) = - targeted_selector(&mut rewritten.lhs).ok_or("left operand is not a selector")?; - let (right, right_restricted) = - targeted_selector(&mut rewritten.rhs).ok_or("right operand is not a selector")?; - if !left.or_matchers.is_empty() || !right.or_matchers.is_empty() { - return Err("selector has an or matcher group"); - } - let is_matching_label = |name: &String| { // PromQL reserves the `__` prefix; none of those names is a join key. !name.starts_with("__") @@ -150,9 +129,32 @@ fn try_propagate( Some(LabelModifier::Include(on)) => on.labels.contains(name), Some(LabelModifier::Exclude(ignoring)) => !ignoring.labels.contains(name), } - && (!left_restricted || crosses_aggregation(name, &left_grouped, &left_ignored)) - && (!right_restricted || crosses_aggregation(name, &right_grouped, &right_ignored)) }; + if let Some(modifier) = &binary.modifier { + let unique_on_matching_labels = |expr: &Expr, tags: &[String]| { + has_unique_aggregate_output(expr) && tags.iter().all(&is_matching_label) + }; + let supported_cardinality = match &modifier.card { + VectorMatchCardinality::OneToOne => true, + VectorMatchCardinality::ManyToOne(_) => { + unique_on_matching_labels(&binary.rhs, right_tags) + } + VectorMatchCardinality::OneToMany(_) => { + unique_on_matching_labels(&binary.lhs, left_tags) + } + VectorMatchCardinality::ManyToMany => false, + }; + if !supported_cardinality { + return Err("one-side uniqueness is not proven"); + } + } + let mut rewritten = binary.clone(); + let left = selector_matchers(&mut rewritten.lhs).ok_or("left operand is not a selector")?; + let right = selector_matchers(&mut rewritten.rhs).ok_or("right operand is not a selector")?; + if !left.or_matchers.is_empty() || !right.or_matchers.is_empty() { + return Err("selector has an or matcher group"); + } + let constraints = left .matchers .iter() @@ -163,8 +165,11 @@ fn try_propagate( let mut changed = false; for matcher in constraints { - for target in [&mut left.matchers, &mut right.matchers] { - if !target.contains(&matcher) { + for (target, operand) in [ + (&mut left.matchers, &binary.lhs), + (&mut right.matchers, &binary.rhs), + ] { + if !target.contains(&matcher) && preserves_filter(operand, &matcher.name) { target.push(matcher.clone()); changed = true; } @@ -184,9 +189,7 @@ fn matches_every_value(matcher: &Matcher) -> bool { /// by a grouping label drops exactly the matching output series and leaves the remaining /// aggregated values untouched. /// -/// `topk` and `bottomk` rank across a group and carry the input labels through, so filtering -/// before them changes the candidate set. `count_values` adds an output label that does not exist -/// in its input. +/// Ranking aggregates instead retain input labels. `count_values` synthesizes a label. fn partitions_by_grouping_labels(aggregate: &AggregateExpr) -> bool { matches!( aggregate.op.id(), @@ -202,22 +205,91 @@ fn partitions_by_grouping_labels(aggregate: &AggregateExpr) -> bool { ) } -/// Returns the matchers of the vector selector an operand scans, together with whether the -/// operand is restricted to the grouping labels an aggregation on its path partitions by, or -/// `None` when the operand's output labels are not proven to be that selector's labels. -fn targeted_selector(expr: &mut Expr) -> Option<(&mut Matchers, bool)> { +/// Ranking operators supported by the planner; limitk and limit_ratio are not implemented. +fn ranks_by_grouping_labels(aggregate: &AggregateExpr) -> bool { + matches!(aggregate.op.id(), token::T_TOPK | token::T_BOTTOMK) +} + +/// Whether a partitioning aggregate proves the operand emits at most one series per combination +/// of its output labels; ranking narrows its input, so it inherits the proof. Uniqueness per +/// *match signature* additionally requires every output label to be a matching label, which the +/// caller checks. +fn has_unique_aggregate_output(expr: &Expr) -> bool { match expr { - Expr::VectorSelector(selector) => Some((&mut selector.matchers, false)), - Expr::Paren(paren) => targeted_selector(&mut paren.expr), - Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => { - let (matchers, _) = targeted_selector(&mut aggregate.expr)?; - Some((matchers, aggregate.modifier.is_some())) + Expr::Paren(paren) => has_unique_aggregate_output(&paren.expr), + Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => true, + Expr::Aggregate(aggregate) if ranks_by_grouping_labels(aggregate) => { + has_unique_aggregate_output(&aggregate.expr) } + _ => false, + } +} + +fn vector_operand_is_lhs(binary: &BinaryExpr) -> Option { + if !matches!( + binary.op.id(), + token::T_ADD | token::T_SUB | token::T_MUL | token::T_DIV | token::T_MOD | token::T_POW + ) || binary.modifier.is_some() + { + return None; + } + if matches!(binary.lhs.as_ref(), Expr::NumberLiteral(_)) { + Some(false) + } else if matches!(binary.rhs.as_ref(), Expr::NumberLiteral(_)) { + Some(true) + } else { + None + } +} + +/// Whether adding a matcher on `label` to the scan [`selector_matchers`] finds leaves the +/// operand's output equal to the subset of its unfiltered output that satisfies the matcher. +/// A filter may cross an aggregate only when it removes whole groups: in particular, a label +/// retained by topk is not necessarily one of its partitioning labels. +fn preserves_filter(expr: &Expr, label: &str) -> bool { + match expr { + Expr::VectorSelector(_) => true, + Expr::Paren(paren) => preserves_filter(&paren.expr, label), + Expr::Aggregate(aggregate) => { + let partition_label = match &aggregate.modifier { + Some(LabelModifier::Include(labels)) => labels.labels.iter().any(|x| x == label), + Some(LabelModifier::Exclude(labels)) => labels.labels.iter().all(|x| x != label), + None => false, + }; + partition_label && preserves_filter(&aggregate.expr, label) + } + Expr::Binary(binary) => match vector_operand_is_lhs(binary) { + Some(true) => preserves_filter(&binary.lhs, label), + Some(false) => preserves_filter(&binary.rhs, label), + None => false, + }, + // A rollup emits at most one output series per input series and carries its labels. + Expr::Call(call) if LABEL_PRESERVING_RANGE_FUNCTIONS.contains(&call.func.name) => true, + _ => false, + } +} + +/// Finds the scanned selector through label-preserving operations. Aggregate partitioning +/// is checked separately for each propagated label. +fn selector_matchers(expr: &mut Expr) -> Option<&mut Matchers> { + match expr { + Expr::VectorSelector(selector) => Some(&mut selector.matchers), + Expr::Paren(paren) => selector_matchers(&mut paren.expr), + Expr::Aggregate(aggregate) + if partitions_by_grouping_labels(aggregate) || ranks_by_grouping_labels(aggregate) => + { + selector_matchers(&mut aggregate.expr) + } + Expr::Binary(binary) => match vector_operand_is_lhs(binary) { + Some(true) => selector_matchers(&mut binary.lhs), + Some(false) => selector_matchers(&mut binary.rhs), + None => None, + }, Expr::Call(call) if LABEL_PRESERVING_RANGE_FUNCTIONS.contains(&call.func.name) => { // Every other argument is a scalar, so position does not matter. let matrix = single_matrix_argument(&call.args.args)?; match call.args.args[matrix].as_mut() { - Expr::MatrixSelector(selector) => Some((&mut selector.vs.matchers, false)), + Expr::MatrixSelector(selector) => Some(&mut selector.vs.matchers), _ => None, } } @@ -226,49 +298,6 @@ fn targeted_selector(expr: &mut Expr) -> Option<(&mut Matchers, bool)> { } } -/// Collects the labels every partitioning aggregation on [`targeted_selector`]' path to the -/// scanned selector groups by (`by(...)`, the `LabelModifier::Include` of an aggregation). -fn modifier_label_names(expr: &Expr) -> Vec<&String> { - let mut grouped = Vec::new(); - collect_modifier_label_names(expr, true, &mut grouped); - grouped -} - -/// Collects the labels every partitioning aggregation on [`targeted_selector`]' path to the -/// scanned selector excludes from its grouping labels (`without(...)`, the -/// `LabelModifier::Exclude` of an aggregation). -fn modifier_excluded_label_names(expr: &Expr) -> Vec<&String> { - let mut ignored = Vec::new(); - collect_modifier_label_names(expr, false, &mut ignored); - ignored -} - -fn collect_modifier_label_names<'a>(expr: &'a Expr, included: bool, out: &mut Vec<&'a String>) { - match expr { - Expr::Paren(paren) => collect_modifier_label_names(&paren.expr, included, out), - Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => { - match aggregate.modifier.as_ref() { - Some(LabelModifier::Include(labels)) if included => out.extend(&labels.labels), - Some(LabelModifier::Exclude(labels)) if !included => out.extend(&labels.labels), - _ => (), - } - collect_modifier_label_names(&aggregate.expr, included, out); - } - _ => (), - } -} - -/// Whether an operand an aggregation partitions by `grouped` labels (with `without(...)`, by -/// every label but `ignored`) still carries `name` as an output label, so that a matcher on -/// `name` reaches it on both sides of that aggregation. -fn crosses_aggregation(name: &String, grouped: &[&String], ignored: &[&String]) -> bool { - if grouped.is_empty() { - !ignored.contains(&name) - } else { - grouped.contains(&name) - } -} - fn single_matrix_argument(args: &[Box]) -> Option { let mut found = None; for (index, arg) in args.iter().enumerate() { @@ -322,7 +351,7 @@ mod tests { #[track_caller] fn assert_rewrite(query: &str, expected: &str) { assert_eq!(rewrite(query).unwrap(), parse(expected).unwrap(), "{query}"); - // The rewrite adds every constraint to both sides, so re-running it is a no-op. + // Re-running the rewrite must be a no-op, including one-way propagation. assert!(rewrite(expected).is_none(), "{expected}"); } @@ -490,12 +519,88 @@ mod tests { ); } + #[test] + fn propagates_through_scalar_arithmetic_and_ranking() { + assert_rewrite( + r#"topk(1, a{host="x"}) / on(host) b"#, + r#"topk(1, a{host="x"}) / on(host) b{host="x"}"#, + ); + assert_rewrite( + r#"a / on(host) bottomk(1, b{host="x"})"#, + r#"a{host="x"} / on(host) bottomk(1, b{host="x"})"#, + ); + assert_rewrite( + r#"(8 * rate(a{host="x"}[5m])) / on(host) topk by(host)(1, b)"#, + r#"(8 * rate(a{host="x"}[5m])) / on(host) topk by(host)(1, b{host="x"})"#, + ); + assert_rewrite( + r#"bottomk without(zone)(2, a / 8) / on(host) b{host="x"}"#, + r#"bottomk without(zone)(2, a{host="x"} / 8) / on(host) b{host="x"}"#, + ); + } + + #[test] + fn propagates_grouped_matching_only_with_a_unique_one_side() { + let cases = [ + ( + r#"(8 * rate(a{host="x"}[5m])) / on(host) group_left topk by(host)(1, max by(host)(b))"#, + r#"(8 * rate(a{host="x"}[5m])) / on(host) group_left topk by(host)(1, max by(host)(b{host="x"}))"#, + vec!["host", "zone"], + vec!["host"], + ), + ( + r#"sum by(host)(a) / on(host) group_right b{host!="x"}"#, + r#"sum by(host)(a{host!="x"}) / on(host) group_right b{host!="x"}"#, + vec!["host"], + vec!["host", "zone"], + ), + ( + r#"a{host=""} / on(host) group_left max by(host)(b)"#, + r#"a{host=""} / on(host) group_left max by(host)(b{host=""})"#, + vec!["host", "zone"], + vec!["host"], + ), + ]; + for (query, expected, left_tags, right_tags) in cases { + assert_eq!( + rewrite_with(query, &left_tags, &right_tags).unwrap(), + parse(expected).unwrap(), + "{query}" + ); + assert!(rewrite_with(expected, &left_tags, &right_tags).is_none()); + } + } + + #[test] + fn preserves_ranking_candidates_and_unproven_cardinality() { + for query in [ + r#"a{host="x"} / on(host) group_left max by(host,zone)(b)"#, + r#"a{host="x"} / on(host) group_left topk by(host)(1, b)"#, + r#"a{host="x"} / on(host) group_left topk(1, max by(host)(b))"#, + r#"topk by(zone)(1, a) / on(host) b{host="x"}"#, + r#"topk by(host)(1, topk(1, a)) / on(host) b{host="x"}"#, + r#"topk by(host)(1, sum by(zone)(a)) / on(host) b{host="x"}"#, + r#"(a + vector(8)) / on(host) b{host="x"}"#, + ] { + assert!(rewrite(query).is_none(), "{query}"); + } + // Even with a unique one-side aggregate, global topk must see every host. + assert!( + rewrite_with( + r#"a{host="x"} / on(host) group_left topk(1, max by(host)(b))"#, + &["host", "zone"], + &["host"] + ) + .is_none() + ); + } + #[test] fn leaves_unproven_semantics_unchanged() { for query in [ r#"a or on(host) b{host="x"}"#, r#"a / on(host) group_left b{host="x"}"#, - r#"sum by(host)(a) / on(host) group_right b{host="x"}"#, + r#"sum by(host,zone)(a) / on(host) group_right b{host="x"}"#, r#"label_replace(a,"host","x","zone",".*") / on(host) b{host="x"}"#, r#"a > on(host) b{host="x"}"#, r#"a / on(host) absent_over_time(b{host="x"}[5m])"#, diff --git a/src/query/src/promql/planner/set_operator.rs b/src/query/src/promql/planner/set_operator.rs new file mode 100644 index 00000000000..7a14ff38ad4 --- /dev/null +++ b/src/query/src/promql/planner/set_operator.rs @@ -0,0 +1,833 @@ +// 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. + +//! Planning of the PromQL set operators (`and`, `unless`, and `or`). + +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::sync::Arc; + +use datafusion::logical_expr::{Cast, Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; +use datafusion::scalar::ScalarValue; +use datafusion_common::{NullEquality, TableReference}; +use datatypes::arrow::datatypes::DataType as ArrowDataType; +use promql::extension_plan::UnionDistinctOn; +use promql_parser::parser::token::{self, TokenType}; +use promql_parser::parser::{BinModifier, LabelModifier, VectorMatchCardinality}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; + +use crate::promql::error::{ + ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu, + MultiFieldsNotSupportedSnafu, Result, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, + UnexpectedTokenSnafu, UnsupportedVectorMatchSnafu, +}; +use crate::promql::planner::{ + OR_FLOAT_FIELD_PREFIX, OR_HISTOGRAM_FIELD_PREFIX, PromPlanner, PromPlannerContext, +}; + +impl PromPlanner { + // TODO(ruihang): change function name + #[allow(clippy::too_many_arguments)] + pub(super) fn or_operator( + &mut self, + left: LogicalPlan, + right: LogicalPlan, + left_tag_cols_set: HashSet, + right_tag_cols_set: HashSet, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, + ) -> Result { + let left_is_empty = Self::is_zero_row_empty_relation(&left); + let right_is_empty = Self::is_zero_row_empty_relation(&right); + match (left_is_empty, right_is_empty) { + (true, false) => { + self.ctx = right_context; + return Ok(right); + } + (false, true) => { + self.ctx = left_context; + return Ok(left); + } + (true, true) => { + self.ctx = left_context; + return Ok(left); + } + (false, false) => {} + } + + ensure!( + !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), + UnexpectedPlanExprSnafu { + desc: "OR operator input has zero columns", + } + ); + let left_has_alternative_samples = + Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); + let right_has_alternative_samples = Self::field_columns_are_alternative_samples( + right.schema(), + &right_context.field_columns, + ); + ensure!( + left_context.field_columns.len() == 1 || left_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" + } + ); + ensure!( + right_context.field_columns.len() == 1 || right_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" + } + ); + + // prepare hash sets + let all_tags = left_tag_cols_set + .union(&right_tag_cols_set) + .cloned() + .collect::>(); + let left_qualifier = left.schema().qualified_field(0).0.cloned(); + let right_qualifier = right.schema().qualified_field(0).0.cloned(); + let left_qualifier_string = left_qualifier + .as_ref() + .map(|l| l.to_string()) + .unwrap_or_default(); + let right_qualifier_string = right_qualifier + .as_ref() + .map(|r| r.to_string()) + .unwrap_or_default(); + let left_time_index_column = + left_context + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let right_time_index_column = + right_context + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: right_qualifier_string.clone(), + })?; + let native_histogram_type = Self::native_histogram_arrow_type(); + let is_numeric = |data_type: &ArrowDataType| { + matches!( + data_type, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) + }; + let left_fields = left_context + .field_columns + .iter() + .map(|name| { + left.schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let right_fields = right_context + .field_columns + .iter() + .map(|name| { + right + .schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let left_field = &left_fields[0]; + let right_field = &right_fields[0]; + let left_field_col = &left_field.0; + let right_field_col = &right_field.0; + let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { + fields.iter().all(|(_, _, data_type)| { + is_numeric(data_type) || data_type == &native_histogram_type + }) + }; + let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { + if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>(), + right_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>() + ), + } + .fail(); + } + true + } else { + (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) + || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) + }; + let target_field_type = if mixed_sample_types { + // Mixed vectors use the existing response representation: one nullable float column + // and one nullable native-histogram column. + ArrowDataType::Float64 + } else if left_field.2 == right_field.2 { + left_field.2.clone() + } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { + ArrowDataType::Float64 + } else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_field.2, right_field.2 + ), + } + .fail(); + }; + let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { + let mut reserved_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + for (name, _, _) in left_fields.iter().chain(&right_fields) { + reserved_names.remove(name); + } + reserved_names.extend(all_tags.iter().cloned()); + let unique_name = |prefix: &str, reserved_names: &mut HashSet| { + let mut index = 0; + loop { + let name = format!("{prefix}{index}"); + index += 1; + if reserved_names.insert(name.clone()) { + break name; + } + } + }; + let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); + let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); + (float_field, histogram_field) + } else { + (left_field_col.clone(), String::new()) + }; + let left_tag_types = left_tag_cols_set + .iter() + .map(|label| { + left.schema() + .fields() + .iter() + .find(|field| field.name() == label) + .map(|field| (label.clone(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) + }) + .collect::>>()?; + let right_tag_types = right_tag_cols_set + .iter() + .map(|label| { + right + .schema() + .fields() + .iter() + .find(|field| field.name() == label) + .map(|field| (label.clone(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) + }) + .collect::>>()?; + let mut target_tag_types = HashMap::with_capacity(all_tags.len()); + for label in &all_tags { + let Some(data_type) = + Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label)) + else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR label {label} has incompatible types: {:?} and {:?}", + left_tag_types.get(label), + right_tag_types.get(label) + ), + } + .fail(); + }; + target_tag_types.insert(label.clone(), data_type); + } + let left_has_tsid = left + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); + let right_has_tsid = right + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); + + // step 0: fill all columns in output schema + let mut all_columns_set = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment + // (e.g. `unknown_metric or some_metric`). + if !(left_has_tsid && right_has_tsid) { + all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME); + } + // remove time index column + all_columns_set.remove(&left_time_index_column); + all_columns_set.remove(&right_time_index_column); + if mixed_sample_types { + for (name, _, _) in left_fields.iter().chain(&right_fields) { + all_columns_set.remove(name); + } + all_columns_set.extend(all_tags.iter().cloned()); + all_columns_set.insert(mixed_float_field_col.clone()); + all_columns_set.insert(mixed_histogram_field_col.clone()); + } else if left_field_col != right_field_col { + // remove field column in the right + all_columns_set.remove(right_field_col); + } + let mut all_columns = all_columns_set.into_iter().collect::>(); + // sort to ensure the generated schema is not volatile + all_columns.sort_unstable(); + // use left time index column name as the result time index column name + all_columns.insert(0, left_time_index_column.clone()); + let mut occupied_column_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + + // step 1: align schema using project, fill non-exist columns with null + let aligned_label_expr = |col: &String, source_types: &HashMap| { + let target_type = &target_tag_types[col]; + if let Some(source_type) = source_types.get(col) { + let expr = DfExpr::Column(Column::new(None::, col)); + if source_type == target_type { + expr + } else { + DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone()) + } + } else { + DfExpr::Literal( + Self::string_scalar_value(target_type, None) + .expect("target label type is a string"), + None, + ) + .alias(col.clone()) + } + }; + let null_histogram = + ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; + let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], + output_col: &String| { + if output_col == &mixed_float_field_col { + if let Some((name, qualifier, data_type)) = fields + .iter() + .find(|(_, _, data_type)| is_numeric(data_type)) + { + let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); + if data_type == &ArrowDataType::Float64 { + expr.alias(output_col) + } else { + DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64)) + .alias(output_col) + } + } else { + DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) + } + } else { + fields + .iter() + .find(|(_, _, data_type)| data_type == &native_histogram_type) + .map(|(name, qualifier, _)| { + DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) + }) + .unwrap_or_else(|| { + DfExpr::Literal(null_histogram.clone(), None).alias(output_col) + }) + } + }; + let left_proj_exprs = all_columns.iter().map(|col| { + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&left_fields, col) + } else if !mixed_sample_types + && col == left_field_col + && left_field.2 != target_field_type + { + DfExpr::Cast(Cast::new( + Box::new(DfExpr::Column(Column::new( + left_field.1.clone(), + left_field_col, + ))), + target_field_type.clone(), + )) + .alias(left_field_col.clone()) + } else if target_tag_types.contains_key(col) { + aligned_label_expr(col, &left_tag_types) + } else { + DfExpr::Column(Column::new(None::, col)) + } + }); + let right_time_index_expr = DfExpr::Column(Column::new( + right_qualifier.clone(), + right_time_index_column, + )) + .alias(left_time_index_column.clone()); + // The field column in right side may not have qualifier (it may be removed by join operation), + // so we need to find it from the schema. + // `skip（1)` to skip the time index column + let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { + // expr + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&right_fields, col) + } else if !mixed_sample_types && col == left_field_col { + let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); + if right_field.2 != target_field_type { + DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone())) + .alias(left_field_col.clone()) + } else if left_field_col != right_field_col { + expr.alias(left_field_col.clone()) + } else { + expr + } + } else if target_tag_types.contains_key(col) { + aligned_label_expr(col, &right_tag_types) + } else { + DfExpr::Column(Column::new(None::, col)) + } + }); + let right_proj_exprs = [right_time_index_expr] + .into_iter() + .chain(right_proj_exprs_without_time_index); + + let left_projected = LogicalPlanBuilder::from(left) + .project(left_proj_exprs) + .context(DataFusionPlanningSnafu)? + .alias(left_qualifier_string.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_projected = LogicalPlanBuilder::from(right) + .project(right_proj_exprs) + .context(DataFusionPlanningSnafu)? + .alias(right_qualifier_string.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 2: compute match columns + let mut match_columns = if let Some(modifier) = modifier + && let Some(matching) = &modifier.matching + { + match matching { + // keeps columns mentioned in `on` + LabelModifier::Include(on) => on.labels.clone(), + // removes columns memtioned in `ignoring` + LabelModifier::Exclude(ignoring) => { + let ignoring = ignoring.labels.iter().cloned().collect::>(); + all_tags.difference(&ignoring).cloned().collect() + } + } + } else { + all_tags.iter().cloned().collect() + }; + // sort to ensure the generated plan is not volatile + match_columns.sort_unstable(); + match_columns.dedup(); + occupied_column_names.extend( + left_projected + .schema() + .fields() + .iter() + .chain(right_projected.schema().fields().iter()) + .map(|field| field.name().clone()), + ); + + let visible_schema = left_projected.schema().clone(); + let visible_left_exprs = left_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let visible_right_exprs = right_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let mut left_match_exprs = Vec::with_capacity(match_columns.len()); + let mut right_match_exprs = Vec::with_capacity(match_columns.len()); + let mut next_internal_column = 0; + + for label in &match_columns { + let left_field = if left_tag_cols_set.contains(label) { + Some( + left_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let right_field = if right_tag_cols_set.contains(label) { + Some( + right_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let data_type = match (left_field.as_ref(), right_field.as_ref()) { + (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { + left_type.clone() + } + (Some((_, left_type)), Some((_, right_type))) => { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" + ), + } + .fail(); + } + (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), + (None, None) => ArrowDataType::Utf8, + }; + let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else { + return UnexpectedPlanExprSnafu { + desc: format!("OR match label {label} must be a string"), + } + .fail(); + }; + let internal_name = loop { + let name = format!("__promql_or_match_{next_internal_column}"); + next_internal_column += 1; + if occupied_column_names.insert(name.clone()) { + break name; + } + }; + left_match_exprs.push(Self::normalized_match_key_expr( + label, + left_field, + &value_type, + &internal_name, + )); + right_match_exprs.push(Self::normalized_match_key_expr( + label, + right_field, + &value_type, + &internal_name, + )); + } + + let left_augmented = LogicalPlanBuilder::from(left_projected) + .project(visible_left_exprs.into_iter().chain(left_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_augmented = LogicalPlanBuilder::from(right_projected) + .project(visible_right_exprs.into_iter().chain(right_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 3: build `UnionDistinctOn` with normalized internal match keys. + let visible_field_count = visible_schema.fields().len(); + let compare_key_indices = + (visible_field_count..visible_field_count + match_columns.len()).collect::>(); + let (time_qualifier, _) = visible_schema + .iter() + .find(|(_, field)| field.name() == &left_time_index_column) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let ts_col_idx = left_augmented + .schema() + .iter() + .position(|(qualifier, field)| { + qualifier == time_qualifier && field.name() == &left_time_index_column + }) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let union_distinct_on = UnionDistinctOn::try_new( + left_augmented, + right_augmented, + compare_key_indices, + ts_col_idx, + ) + .context(DataFusionPlanningSnafu)?; + let augmented_result = LogicalPlan::Extension(Extension { + node: Arc::new(union_distinct_on), + }); + let result = LogicalPlanBuilder::from(augmented_result) + .project(visible_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 4: update context + let output_field_col = left_field_col.clone(); + let mut output_context = left_context; + let mut visible_tags = all_tags.into_iter().collect::>(); + visible_tags.sort_unstable(); + output_context.time_index_column = Some(left_time_index_column); + output_context.tag_columns = visible_tags; + output_context.field_columns = if mixed_sample_types { + vec![mixed_float_field_col, mixed_histogram_field_col] + } else { + vec![output_field_col] + }; + output_context.use_tsid = left_has_tsid && right_has_tsid; + self.ctx = output_context; + + Ok(result) + } + + /// Build a set operator (AND/OR/UNLESS) + pub(super) fn set_op_on_non_field_columns( + &mut self, + mut left: LogicalPlan, + mut right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + op: TokenType, + modifier: &Option, + ) -> Result { + let left_tag_col_set = left_context + .tag_columns + .iter() + .cloned() + .collect::>(); + let right_tag_col_set = right_context + .tag_columns + .iter() + .cloned() + .collect::>(); + + if matches!(op.id(), token::T_LOR) { + return self.or_operator( + left, + right, + left_tag_col_set, + right_tag_col_set, + left_context, + right_context, + modifier, + ); + } + + if let Some(modifier) = modifier { + ensure!( + matches!( + modifier.card, + VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany + ), + UnsupportedVectorMatchSnafu { + name: modifier.card.clone(), + }, + ); + } + + let output_context = left_context.clone(); + let visible_left_schema = left.schema().clone(); + let mut left_context = left_context; + let mut right_context = right_context; + let added_marker_to_left = if Self::only_temporality_match_label_mismatches( + &left_context, + &right_context, + modifier, + ) { + let aligned = Self::align_temporality_match_column( + left, + right, + &mut left_context, + &mut right_context, + )?; + left = aligned.0; + right = aligned.1; + aligned.2 + } else { + false + }; + + let mut left_tag_col_set = left_context + .tag_columns + .iter() + .cloned() + .collect::>(); + let mut right_tag_col_set = right_context + .tag_columns + .iter() + .cloned() + .collect::>(); + if let Some(matching) = modifier + .as_ref() + .and_then(|modifier| modifier.matching.as_ref()) + { + match matching { + LabelModifier::Include(on) => { + let mask = on.labels.iter().cloned().collect::>(); + left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect(); + right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect(); + } + LabelModifier::Exclude(ignoring) => { + for label in &ignoring.labels { + let _ = left_tag_col_set.remove(label); + let _ = right_tag_col_set.remove(label); + } + } + } + } + ensure!( + left_tag_col_set == right_tag_col_set, + CombineTableColumnMismatchSnafu { + left: left_tag_col_set.iter().cloned().collect::>(), + right: right_tag_col_set.iter().cloned().collect::>(), + } + ); + + let left_time_index = left_context.time_index_column.clone().unwrap(); + let right_time_index = right_context.time_index_column.clone().unwrap(); + + // alias right time index column if necessary + if left_context.time_index_column != right_context.time_index_column { + let right_project_exprs = right + .schema() + .fields() + .iter() + .map(|field| { + if field.name() == &right_time_index { + DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index) + } else { + DfExpr::Column(Column::from_name(field.name())) + } + }) + .collect::>(); + + right = LogicalPlanBuilder::from(right) + .project(right_project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + } + + let join_keys = left_tag_col_set + .into_iter() + .chain([left_time_index]) + .map(Column::from_name) + .collect::>(); + + ensure!( + left_context.field_columns.len() == 1 + || Self::field_columns_are_alternative_samples( + left.schema(), + &left_context.field_columns, + ), + MultiFieldsNotSupportedSnafu { + operator: "AND/UNLESS operator" + } + ); + // Generate join plan. + // All set operations in PromQL are "distinct" + let result = match op.id() { + token::T_LAND => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftSemi, + (join_keys.clone(), join_keys), + None, + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LUNLESS => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftAnti, + (join_keys.clone(), join_keys), + None, + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LOR => { + // OR is handled at the beginning of this function, as it cannot + // be expressed using JOIN like AND and UNLESS. + unreachable!() + } + _ => UnexpectedTokenSnafu { token: op }.fail(), + }?; + let result = if added_marker_to_left { + LogicalPlanBuilder::from(result) + .project(visible_left_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + result + }; + + // AND/UNLESS preserve the complete left operand's visible columns and values; encoded + // markers are decoded. + self.ctx = output_context; + Ok(result) + } +} diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs new file mode 100644 index 00000000000..19de1098288 --- /dev/null +++ b/src/query/src/promql/planner/test.rs @@ -0,0 +1,8902 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use catalog::RegisterTableRequest; +use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; +use common_base::Plugins; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_query::native_histogram::{ + CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, +}; +use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use common_query::test_util::DummyDecoder; +use common_recordbatch::RecordBatch as GreptimeRecordBatch; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::{MemTable, provider_as_source}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::{Extension, col}; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use promql::extension_plan::HistogramFold; +use promql_parser::label::Labels; +use promql_parser::parser; +use session::context::QueryContext; +use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; +use table::Table; +use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; +use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; + +use super::*; +use crate::QueryEngineContext; +use crate::options::QueryOptions; +use crate::parser::QueryLanguageParser; +use crate::query_engine::DefaultSerializer; + +mod delta; + +fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { + if let LogicalPlan::Extension(Extension { node }) = plan + && let Some(instant_manipulate) = node.as_any().downcast_ref::() + { + return Some(instant_manipulate); + } + + plan.inputs().into_iter().find_map(find_instant_manipulate) +} + +fn build_query_engine_state() -> QueryEngineState { + QueryEngineState::new( + new_memory_catalog_manager().unwrap(), + None, + None, + None, + None, + None, + false, + Plugins::default(), + QueryOptions::default(), + ) +} + +#[test] +fn common_label_type_preserves_only_shared_dictionary_encoding() { + let dictionary = ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Utf8), + ); + let other_dictionary = ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ); + + assert_eq!( + Some(dictionary.clone()), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), None) + ); +} + +async fn build_optimized_promql_plan( + table_provider: DfTableSourceProvider, + eval_stmt: &EvalStmt, +) -> LogicalPlan { + let state = build_query_engine_state(); + let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state) + .await + .unwrap(); + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + state + .optimize_by_extension_rules(raw_plan, &context) + .unwrap() +} + +async fn build_optimized_tsid_plan( + query: &str, + num_tag: usize, + num_field: usize, + end_secs: u64, + lookback_secs: u64, +) -> String { + let eval_stmt = EvalStmt { + expr: parser::parse(query).unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(end_secs)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(lookback_secs), + }; + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + num_tag, + num_field, + ) + .await; + + build_optimized_promql_plan(table_provider, &eval_stmt) + .await + .display_indent_schema() + .to_string() +} + +async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) { + let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; + + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0")); + assert!(plan_str.contains("Distinct:")); + assert!(plan_str.contains(expected_outer_agg), "{plan_str}"); + assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); +} + +async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) { + let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await; + assert!(!plan_str.contains("Distinct:"), "{plan_str}"); +} + +fn build_eval_stmt(expr: &str) -> EvalStmt { + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } +} + +enum DirectOrValue { + Float64(f64), + Int64(i64), + NativeHistogram(NativeHistogram), + Utf8(&'static str), +} + +impl DirectOrValue { + fn data_type(&self) -> ArrowDataType { + match self { + Self::Float64(_) => ArrowDataType::Float64, + Self::Int64(_) => ArrowDataType::Int64, + Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), + Self::Utf8(_) => ArrowDataType::Utf8, + } + } + fn array(&self) -> Arc { + match self { + Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), + Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), + Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), + } + } +} + +fn direct_or_histogram() -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum: 1.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![], + positive_spans: vec![], + negative_spans: vec![], + count: 1.0, + zero_count: 1.0, + positive_buckets: vec![], + negative_buckets: vec![], + } +} + +fn operator_metric_table( + name: &str, + table_id: u32, + tag: &str, + le: Option<&str>, + value: DirectOrValue, +) -> table::TableRef { + let value_type = match &value { + DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), + DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), + DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), + DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), + }; + let tag_count = 1 + usize::from(le.is_some()); + let mut columns = vec![ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if le.is_some() { + columns.push(ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.extend([ + ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("v".to_string(), value_type, true), + ]); + let schema = Arc::new(Schema::new(columns)); + let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; + if let Some(le) = le { + arrays.push(Arc::new(StringArray::from(vec![le]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, + value.array(), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + name, + GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), + table_id, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let value_index = tag_count + 1; + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![value_index]) + .next_column_id((value_index + 1) as u32) + .build() + .unwrap(); + let info = Arc::new( + TableInfoBuilder::default() + .table_id(table_id) + .name(name) + .meta(meta) + .build() + .unwrap(), + ); + Arc::new(Table::new( + info, + FilterPushDownType::Unsupported, + backing.data_source(), + )) +} + +fn operator_table_provider() -> DfTableSourceProvider { + let catalog = MemoryCatalogManager::with_default_setup(); + let tables = [ + operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), + operator_metric_table( + "lh", + 2_002, + "b", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), + operator_metric_table( + "rh", + 2_004, + "a", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), + operator_metric_table( + "bad_classic", + 2_006, + "d", + Some("broken"), + DirectOrValue::Float64(1.0), + ), + operator_metric_table( + "bad_native", + 2_007, + "d", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + ]; + for table in tables { + let info = table.table_info(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: info.name.clone(), + table_id: info.ident.table_id, + table, + }) + .unwrap(); + } + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn operator_eval_stmt(expr: &str) -> EvalStmt { + let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: time, + end: time, + interval: Duration::from_secs(1), + lookback_delta: Duration::from_secs(5), + } +} + +struct DirectOrSource { + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, +} + +fn source( + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, +) -> DirectOrSource { + DirectOrSource { + name, + empty, + timestamp, + tags, + value, + } +} + +fn tagged_source( + name: &'static str, + empty: bool, + tag: (&'static str, Option<&'static str>), + value: DirectOrValue, +) -> DirectOrSource { + source(name, empty, 1, vec![("job", Some("job")), tag], value) +} + +fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { + source(name, true, 1, vec![("job", Some("job"))], value) +} + +fn table(source: &DirectOrSource) -> Arc { + let mut fields = vec![Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )]; + fields.extend( + source + .tags + .iter() + .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), + ); + fields.push(Field::new("v", source.value.data_type(), true)); + let schema = Arc::new(ArrowSchema::new(fields)); + let partitions = if source.empty { + vec![vec![]] + } else { + let mut columns: Vec> = + vec![Arc::new(TimestampMillisecondArray::from(vec![ + source.timestamp, + ]))]; + columns.extend( + source + .tags + .iter() + .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), + ); + columns.push(source.value.array()); + vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] + }; + Arc::new(MemTable::try_new(schema, partitions).unwrap()) +} + +fn scan(source: &DirectOrSource) -> LogicalPlan { + LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) + .unwrap() + .build() + .unwrap() +} + +fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { + PromPlannerContext { + table_name: Some(qualifier.to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![field.to_string()], + tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), + ..Default::default() + } +} + +fn or_modifier(expr: &str) -> Option { + let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { + unreachable!() + }; + expr.modifier +} + +async fn plan_direct_or( + left: LogicalPlan, + right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, +) -> LogicalPlan { + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + planner + .or_operator( + left, + right, + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + modifier, + ) + .unwrap() +} + +async fn execute(plan: LogicalPlan, state: &QueryEngineState) -> (LogicalPlan, Vec) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + let batches = datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + (optimized, batches) +} + +async fn run( + left: &DirectOrSource, + right: &DirectOrSource, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, +) -> (LogicalPlan, Vec) { + let plan = plan_direct_or( + scan(left), + scan(right), + left_context, + right_context, + modifier, + ) + .await; + execute(plan, &build_query_engine_state()).await +} + +async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { + let sample = |histogram: bool| { + if histogram { + DirectOrValue::NativeHistogram(direct_or_histogram()) + } else { + DirectOrValue::Float64(1.25) + } + }; + let left = tagged_source( + "lhs", + false, + ( + "k", + Some(if histogram_on_left { + "histogram" + } else { + "float" + }), + ), + sample(histogram_on_left), + ); + let right = tagged_source( + "rhs", + false, + ( + "k", + Some(if histogram_on_left { + "float" + } else { + "histogram" + }), + ), + sample(!histogram_on_left), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job", "k"], "v"); + let right_context = direct_or_context("rhs", &["job", "k"], "v"); + let plan = planner + .or_operator( + scan(&left), + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + (planner, plan) +} + +async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { + let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); + let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); + let row_count = histograms.len() + 1; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new("k", ArrowDataType::Utf8, false), + Field::new(&float_field, ArrowDataType::Float64, true), + Field::new( + &histogram_field, + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let mut histogram_values = Vec::with_capacity(row_count); + histogram_values.push(None); + histogram_values.extend(histograms.into_iter().map(Some)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), + Arc::new(StringArray::from_iter_values( + (0..row_count).map(|row| format!("kind_{row}")), + )), + Arc::new(Float64Array::from_iter( + (0..row_count).map(|row| (row == 0).then_some(1.25)), + )), + build_histogram_array(&histogram_values), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let planner = PromPlanner { + table_provider, + ctx: PromPlannerContext { + table_name: Some("mixed".to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![float_field, histogram_field], + tag_columns: vec!["k".to_string()], + ..Default::default() + }, + promql_annotations: None, + }; + (planner, plan) +} + +fn assert_no_internal_or_keys(schema: &DFSchema) { + assert!( + schema + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")), + "{schema:?}" + ); +} + +fn values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect() +} + +fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = datafusion::arrow::compute::cast( + batch.column_by_name(column).unwrap(), + &ArrowDataType::Float64, + ) + .unwrap(); + values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .collect::>() + }) + .collect() +} + +fn histograms(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).filter_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .collect() +} + +fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name("v") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column_by_name("k") + .map(|column| column.as_any().downcast_ref::().unwrap()); + (0..batch.num_rows()).map(move |i| { + ( + values.value(i), + labels.and_then(|labels| { + (!labels.is_null(i)).then(|| labels.value(i).to_string()) + }), + ) + }) + }) + .collect::>(); + rows.sort_by(|left, right| left.0.total_cmp(&right.0)); + rows +} + +fn matrix_source( + name: &'static str, + k: Option>, + timestamp: i64, + value: f64, +) -> DirectOrSource { + let mut tags = vec![("job", Some("job"))]; + if let Some(k) = k { + tags.push(("k", k)); + } + source(name, false, timestamp, tags, DirectOrValue::Float64(value)) +} + +fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { + direct_or_context( + name, + if k.is_some() { &["job", "k"] } else { &["job"] }, + "v", + ) +} + +async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "non_existent_histogram_bucket".to_string(), + ), + (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()), + ], + &["pod", "instance"], + ) + .await +} + +fn assert_normal_metric_schema(plan: &LogicalPlan) { + let fields = plan.schema().fields(); + assert_eq!(fields.len(), 4, "{fields:?}"); + assert!( + fields.iter().any(|field| field.name() == "pod"), + "{fields:?}" + ); + assert!( + fields.iter().any(|field| field.name() == "instance"), + "{fields:?}" + ); + assert!( + fields + .iter() + .any(|field| field.name() == greptime_timestamp()), + "{fields:?}" + ); + assert!( + fields.iter().any(|field| { + field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64 + }), + "{fields:?}" + ); +} + +async fn build_test_table_provider_with_distinct_tags( + table_tags: &[(&str, &[&str])], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (table_name, tags) in table_tags { + let mut columns = tags + .iter() + .map(|tag| { + ColumnSchema::new( + (*tag).to_string(), + ConcreteDataType::string_datatype(), + false, + ) + }) + .collect::>(); + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(columns))) + .primary_key_indices((0..tags.len()).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name((*table_name).to_string()) + .meta(table_meta) + .build() + .unwrap(); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: (*table_name).to_string(), + table_id: 1024, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn contains_histogram_fold(plan: &LogicalPlan) -> bool { + matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::()) + || plan.inputs().into_iter().any(contains_histogram_fold) +} + +async fn build_set_op_context_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_distinct_tags(&[ + ("bucket_metric", &["job", "le"]), + ("normal_metric", &["job"]), + ("fallback_metric", &["instance"]), + ]) + .await +} + +async fn build_or_context_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_distinct_tags(&[ + ("normal_metric", &["job"]), + ("other_metric", &["instance"]), + ("non_hist_metric", &["instance"]), + ]) + .await +} + +async fn optimize_and_create_physical_plan( + state: &QueryEngineState, + plan: LogicalPlan, +) -> ( + LogicalPlan, + Arc, +) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + (optimized, physical) +} + +async fn build_test_table_provider( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (schema_name, table_name) in table_name_tuples { + let mut columns = vec![]; + for i in 0..num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..num_tag).collect()) + .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id: 1024, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { + build_test_native_histogram_table_provider_with_marker(table_name, false).await +} + +async fn build_test_native_histogram_table_provider_with_marker( + table_name: &str, + temporality_marker: bool, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let mut columns = vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ]; + if temporality_marker { + columns.push(ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + )); + } + let tag_count = columns.len(); + columns.extend([ + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ]); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![tag_count + 1]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + "native_histogram_2".to_string(), + native_histogram_value_type().clone(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2, 3]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_mixed_native_histogram_table_provider( + table_name: &str, +) -> DfTableSourceProvider { + build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await +} + +async fn build_test_mixed_native_histogram_table_provider_with_marker( + table_name: &str, + temporality_marker: bool, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let mut columns = vec![ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if temporality_marker { + columns.push(ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + )); + } + let tag_count = columns.len(); + columns.extend([ + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![tag_count + 1, tag_count + 2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let mut arrays: Vec> = + vec![Arc::new(StringArray::from(vec!["float", "histogram"]))]; + if temporality_marker { + arrays.push(Arc::new(StringArray::from(vec![ + Some(GREPTIME_TEMPORALITY_DELTA), + Some(GREPTIME_TEMPORALITY_DELTA), + ]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc, + build_histogram_array(&[None, Some(direct_or_histogram())]), + Arc::new(Float64Array::from(vec![Some(2.0), None])), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 1024, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn classic_and_native_histogram_table_provider( + native_tag: &str, + native_le: Option<&str>, + native_histogram: NativeHistogram, +) -> DfTableSourceProvider { + let table_name = "mixed_histogram"; + let catalog = MemoryCatalogManager::with_default_setup(); + let schema = Arc::new(Schema::new(vec![ + // A dotted tag name guards the mixed histogram_quantile projection + // against qualified-name parsing (#9390). + ColumnSchema::new( + "service.name".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3, 4]) + .next_column_id(5) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec![ + "classic", "classic", native_tag, "classic", "classic", native_tag, + ])), + Arc::new(StringArray::from(vec![ + Some("1"), + Some("+Inf"), + native_le, + Some("1"), + Some("+Inf"), + native_le, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, + ])), + build_histogram_array(&[ + None, + None, + Some(native_histogram.clone()), + None, + None, + Some(native_histogram), + ]), + Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + None, + Some(2.0), + Some(4.0), + None, + ])), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 2_200, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 2_200, + table, + }) + .unwrap(); + + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_table_provider_with_tsid( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, +) -> DfTableSourceProvider { + let table_specs = table_name_tuples + .iter() + .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field)) + .collect::>(); + build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await +} + +async fn build_test_table_provider_with_tsid_fields( + table_specs: &[((String, String), usize)], + num_tag: usize, +) -> DfTableSourceProvider { + let table_specs = table_specs + .iter() + .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field)) + .collect::>(); + build_test_table_provider_with_tsid_tag_fields(&table_specs).await +} + +async fn build_test_table_provider_with_tsid_tag_fields( + table_specs: &[((String, String), usize, usize)], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + + let physical_table_name = "phy"; + let physical_table_id = 999u32; + let physical_num_tag = table_specs + .iter() + .map(|(_, num_tag, _)| *num_tag) + .max() + .unwrap_or(0); + let physical_num_field = table_specs + .iter() + .map(|(_, _, num_field)| *num_field) + .max() + .unwrap_or(0); + + // Register a metric engine physical table with internal columns. + { + let mut columns = vec![ + ColumnSchema::new( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + ConcreteDataType::uint32_datatype(), + false, + ), + ColumnSchema::new( + DATA_SCHEMA_TSID_COLUMN_NAME.to_string(), + ConcreteDataType::uint64_datatype(), + false, + ), + ]; + for i in 0..physical_num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..physical_num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + + let schema = Arc::new(Schema::new(columns)); + let primary_key_indices = (0..(2 + physical_num_tag)).collect::>(); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(primary_key_indices) + .value_indices( + (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(), + ) + .engine(METRIC_ENGINE_NAME.to_string()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(physical_table_id) + .name(physical_table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: physical_table_name.to_string(), + table_id: physical_table_id, + table, + }) + .is_ok() + ); + } + + // Register metric engine logical tables without `__tsid`, referencing the physical table. + for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate() { + let mut columns = vec![]; + for i in 0..*num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..*num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + + let schema = Arc::new(Schema::new(columns)); + let mut options = table::requests::TableOptions::default(); + options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + physical_table_name.to_string(), + ); + let table_id = 1024u32 + idx as u32; + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..*num_tag).collect()) + .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect()) + .engine(METRIC_ENGINE_NAME.to_string()) + .options(options) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(table_id) + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_table_provider_with_fields( + table_name_tuples: &[(String, String)], + tags: &[&str], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (schema_name, table_name) in table_name_tuples { + let mut columns = vec![]; + let num_tag = tags.len(); + for tag in tags { + columns.push(ColumnSchema::new( + tag.to_string(), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..num_tag).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id: 1024, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +// { +// input: `abs(some_metric{foo!="bar"})`, +// expected: &Call{ +// Func: MustGetFunction("abs"), +// Args: Expressions{ +// &VectorSelector{ +// Name: "some_metric", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// }, +// }, +// }, +async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) { + let prom_expr = parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, 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]\ + \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]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]" + ).replace("TEMPLATE", plan_name); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn single_abs() { + do_single_instant_function_call("abs", "abs").await; +} + +#[tokio::test] +#[should_panic] +async fn single_absent() { + do_single_instant_function_call("absent", "").await; +} + +#[tokio::test] +async fn single_ceil() { + do_single_instant_function_call("ceil", "ceil").await; +} + +#[tokio::test] +async fn single_exp() { + do_single_instant_function_call("exp", "exp").await; +} + +#[tokio::test] +async fn single_ln() { + do_single_instant_function_call("ln", "ln").await; +} + +#[tokio::test] +async fn single_log2() { + do_single_instant_function_call("log2", "log2").await; +} + +#[tokio::test] +async fn single_log10() { + do_single_instant_function_call("log10", "log10").await; +} + +#[tokio::test] +#[should_panic] +async fn single_scalar() { + do_single_instant_function_call("scalar", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sgn() { + do_single_instant_function_call("sgn", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sort() { + do_single_instant_function_call("sort", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sort_desc() { + do_single_instant_function_call("sort_desc", "").await; +} + +#[tokio::test] +async fn single_sqrt() { + do_single_instant_function_call("sqrt", "sqrt").await; +} + +#[tokio::test] +async fn single_timestamp_plan_preserves_source_value() { + let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \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 Filter: some_metric.field_0 IS NOT NULL [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \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 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]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn single_acos() { + do_single_instant_function_call("acos", "acos").await; +} + +#[tokio::test] +#[should_panic] +async fn single_acosh() { + do_single_instant_function_call("acosh", "").await; +} + +#[tokio::test] +async fn single_asin() { + do_single_instant_function_call("asin", "asin").await; +} + +#[tokio::test] +#[should_panic] +async fn single_asinh() { + do_single_instant_function_call("asinh", "").await; +} + +#[tokio::test] +async fn single_atan() { + do_single_instant_function_call("atan", "atan").await; +} + +#[tokio::test] +#[should_panic] +async fn single_atanh() { + do_single_instant_function_call("atanh", "").await; +} + +#[tokio::test] +async fn single_cos() { + do_single_instant_function_call("cos", "cos").await; +} + +#[tokio::test] +#[should_panic] +async fn single_cosh() { + do_single_instant_function_call("cosh", "").await; +} + +#[tokio::test] +async fn single_sin() { + do_single_instant_function_call("sin", "sin").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sinh() { + do_single_instant_function_call("sinh", "").await; +} + +#[tokio::test] +async fn single_tan() { + do_single_instant_function_call("tan", "tan").await; +} + +#[tokio::test] +#[should_panic] +async fn single_tanh() { + do_single_instant_function_call("tanh", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_deg() { + do_single_instant_function_call("deg", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_rad() { + do_single_instant_function_call("rad", "").await; +} + +// { +// input: "avg by (foo)(some_metric)", +// expected: &AggregateExpr{ +// Op: AVG, +// Expr: &VectorSelector{ +// Name: "some_metric", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// PosRange: PositionRange{ +// Start: 13, +// End: 24, +// }, +// }, +// Grouping: []string{"foo"}, +// PosRange: PositionRange{ +// Start: 0, +// End: 25, +// }, +// }, +// }, +async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { + let prom_expr = parser::parse(&format!( + "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})", + )) + .unwrap(); + let mut eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // test group by + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 2, + 2, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected_no_without = String::from( + "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1: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, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" + ).replace("TEMPLATE", plan_name); + assert_eq!( + plan.display_indent_schema().to_string(), + expected_no_without + ); + + // test group without + if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr { + *modifier = Some(LabelModifier::Exclude(Labels { + labels: vec![String::from("tag_1")].into_iter().collect(), + })); + } + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 2, + 2, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected_without = String::from( + "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1: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, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" + ).replace("TEMPLATE", plan_name); + assert_eq!(plan.display_indent_schema().to_string(), expected_without); +} + +#[tokio::test] +async fn aggregate_sum() { + do_aggregate_expr_plan("sum", "sum").await; +} + +#[tokio::test] +async fn tsid_is_used_for_series_divide_when_available() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + assert!(plan_str.contains("__tsid ASC NULLS FIRST")); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); + + let manipulate = find_instant_manipulate(&plan).unwrap(); + let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new( + 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")); +} + +async fn build_at_modifier_plan(query: &str, start_secs: u64, end_secs: u64) -> LogicalPlan { + let eval_stmt = build_at_modifier_eval_stmt(query, start_secs, end_secs); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap() +} + +fn build_at_modifier_eval_stmt(query: &str, start_secs: u64, end_secs: u64) -> EvalStmt { + EvalStmt { + expr: parser::parse(query).unwrap(), + start: UNIX_EPOCH + .checked_add(Duration::from_secs(start_secs)) + .unwrap(), + end: UNIX_EPOCH + .checked_add(Duration::from_secs(end_secs)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } +} + +/// Every selector of an `@` anchored query must scan around the anchor only, instead of the +/// whole evaluation range. +#[tokio::test] +async fn at_modifier_anchors_selector_scan_window() { + // `@ 100` anchors at t=100s; the lookback delta is 1s, so the scan covers (99s, 100s]. + let plan = build_at_modifier_plan("some_metric @ 100", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(99001, None) AND some_metric.timestamp <= TimestampMillisecond(100000, None)" + ), + "{plan_str}" + ); + // The result is reported at the evaluation timestamps, not at the anchor. + assert!(plan_str.contains("range=[0..1000000]"), "{plan_str}"); + + // `@ start()` / `@ end()` resolve to the evaluation range of the statement. + let plan = build_at_modifier_plan("some_metric @ start()", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(0, None)" + ), + "{plan_str}" + ); + + let plan = build_at_modifier_plan("some_metric @ end()", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(999001, None) AND some_metric.timestamp <= TimestampMillisecond(1000000, None)" + ), + "{plan_str}" + ); + + // `offset` moves the anchor backwards and is not applied twice. + let plan = build_at_modifier_plan("some_metric @ 200 offset 50s", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(149001, None) AND some_metric.timestamp <= TimestampMillisecond(150000, None)" + ), + "{plan_str}" + ); + + // A timestamp before the Unix epoch is accepted, as in Prometheus. + let plan = build_at_modifier_plan("some_metric @ -1", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(-1999, None) AND some_metric.timestamp <= TimestampMillisecond(-1000, None)" + ), + "{plan_str}" + ); +} + +/// A call over one range selector anchored by `@` is evaluated once, at the start of the +/// evaluation, and its result is reported at every step: the window is folded around the anchor +/// instead of following the outer evaluation grid. This is the planner's counterpart of +/// Prometheus' `StepInvariantExpr` wrapper; see [`PromPlanner::promotes_anchored_range_call`]. +#[tokio::test] +async fn at_modifier_promotes_anchored_range_call() { + let plan = build_at_modifier_plan("rate(some_metric[5m] @ 300)", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The scan is limited to the anchored window (offset by `eval_start - anchor`). + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(1, None) AND some_metric.timestamp <= TimestampMillisecond(300000, None)" + ), + "{plan_str}" + ); + // A single fold, at the anchor... + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + // ... never a fold per step of the outer grid. + assert!( + !plan_str.contains("PromRangeManipulate: req range=[0..1000000]"), + "{plan_str}" + ); + // A single replay of the function result over the whole grid... + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{plan_str}" + ); + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000001]"), + "{plan_str}" + ); + // ... so `rate` itself is evaluated below that replay node, on the single evaluation + // instant of the anchored subtree, instead of once per step. + let replay = plan_str + .find("PromInstantManipulate") + .expect("instant manipulate node"); + let rate = plan_str.find("prom_rate(").expect("rate projection"); + assert!( + replay < rate, + "`rate` must be evaluated below the replay node:\n{plan_str}" + ); +} + +/// Parentheses around the range argument are transparent: `rate((some_metric[5m] @ 300))` gets +/// the same fixed-window promotion as `rate(some_metric[5m] @ 300)`, with the anchored window +/// folded once and its result replayed at every step of the grid. Only that one argument is +/// looked through, so a parenthesis above the call promotes no operator of its own and a +/// parenthesized subtree is planned exactly like the bare one. +#[tokio::test] +async fn at_modifier_promotes_parenthesized_range_argument() { + // Each form is planned exactly like its unparenthesized counterpart: parentheses below the + // call are transparent, a parenthesis around the call adds nothing, and a parenthesis above + // it does not widen the promotion. + for (query, plain) in [ + ( + "rate((some_metric[5m] @ 300))", + "rate(some_metric[5m] @ 300)", + ), + ( + "rate(((some_metric[5m] @ 300)))", + "rate(some_metric[5m] @ 300)", + ), + ( + "(rate(some_metric[5m] @ 300))", + "rate(some_metric[5m] @ 300)", + ), + ( + "abs((rate(some_metric[5m] @ 300)))", + "abs(rate(some_metric[5m] @ 300))", + ), + ] { + assert_eq!( + build_at_modifier_plan(query, 0, 1000) + .await + .display_indent_schema() + .to_string(), + build_at_modifier_plan(plain, 0, 1000) + .await + .display_indent_schema() + .to_string(), + "`{query}` must be planned like `{plain}`" + ); + } + + // The parentheses do not push the enclosing operator into the promotion either: `abs` stays + // above the replay of the promoted call, exactly as it does without them. The promotion + // itself (one anchored fold, one replay, `rate` below it) is asserted for the + // unparenthesized form by `at_modifier_promotes_anchored_range_call`, and the form above is + // planned identically to it. + let plan_str = build_at_modifier_plan("abs((rate(some_metric[5m] @ 300)))", 0, 1000) + .await + .display_indent_schema() + .to_string(); + let replay = plan_str + .find("PromInstantManipulate") + .expect("instant manipulate node"); + assert!( + plan_str.find("abs(").expect("`abs` projection") < replay, + "`abs` must be evaluated above the replay of the promoted call:\n{plan_str}" + ); +} + +/// `@ start()` and `@ end()` are fixed anchors for the whole statement, so a call using them is +/// promoted as well. +#[tokio::test] +async fn at_modifier_promotes_start_and_end_anchored_call() { + for query in [ + "rate(some_metric[5m] @ start())", + "rate(some_metric[5m] @ end())", + "max_over_time(some_metric[5m] @ end())", + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000001]"), + "{query}:\n{plan_str}" + ); + } + + // Only the call itself is promoted: a call above it (`abs`) is planned as usual and + // evaluated at every step over the replayed result of the promoted call. The window is + // still folded once, around the anchor. + let plan = build_at_modifier_plan("abs(max_over_time(some_metric[5m] @ end()))", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + let abs = plan_str.find("abs(").expect("`abs` projection"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + abs < replay, + "`abs` must be evaluated above the replay of the promoted call:\n{plan_str}" + ); + // The window is the only one folded once, and it sits below the replay: the promoted call + // feeds the grid from there. + let fold = plan_str.find("PromRangeManipulate").expect("range fold"); + assert!( + replay < fold, + "the anchored window must be folded below the replay:\n{plan_str}" + ); + + // An aggregation above the promoted call stays above it as well: `sum` aggregates the + // replayed per-series rows at every step, instead of aggregating the single anchored + // instant and replaying the aggregation — which would also have to replay rows of several + // groups through the one-series-per-batch `InstantManipulate`. + let plan = build_at_modifier_plan("sum(rate(some_metric[5m] @ start()))", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + // A single replay, and it sits below the aggregation node: `sum` aggregates the replayed + // per-series rows at every step. + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{plan_str}" + ); + let aggregate = plan_str.find("Aggregate:").expect("aggregate node"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + aggregate < replay, + "the aggregation must stay above the replay of the promoted call:\n{plan_str}" + ); + + // Without `@` nothing is promoted: the function keeps folding one window per step. + let plan = build_at_modifier_plan("rate(some_metric[5m])", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("PromRangeManipulate: req range=[0..1000000]"), + "{plan_str}" + ); + assert!(!plan_str.contains("lookback=[1000001]"), "{plan_str}"); +} + +/// A call or a unary operator above the anchored range call is planned as usual: the inner range +/// call is promoted on its own (it is the direct call over the anchored range selector), and the +/// operator above it is evaluated at every step over the replayed result. +#[tokio::test] +async fn at_modifier_promotes_inner_range_call_below_wrappers() { + for (query, wrapper) in [ + ("abs(rate(some_metric[5m] @ 300))", "abs(prom_rate("), + ("-rate(some_metric[5m] @ 300)", "(- prom_rate("), + ( + "abs(max_over_time(some_metric[5m] @ 300))", + "abs(prom_max_over_time(", + ), + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The anchored window is folded once, and the wrapper sits above its replay. + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{query}:\n{plan_str}" + ); + let wrapper = plan_str + .find(wrapper) + .unwrap_or_else(|| panic!("no `{wrapper}` projection in:\n{plan_str}")); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + wrapper < replay, + "the wrapper must be evaluated above the replay of the range call:\n{plan_str}" + ); + } +} + +/// `anchored + plain`: only the binary operand that is a call over the anchored range selector +/// is promoted, and the plain side keeps following the evaluation step. +#[tokio::test] +async fn at_modifier_promotes_only_anchored_binary_operand() { + let plan = build_at_modifier_plan("rate(some_metric[5m] @ 300) + some_metric", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The anchored operand is folded once and its result replayed over the whole grid. + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000001]"), + "{plan_str}" + ); + // The plain operand still selects one sample per step with the default lookback. + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000]"), + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); + + // Both operands anchored: the binary expression is not promoted as a whole, because a join + // emits the rows of several series in shared batches and the replay needs one series per + // batch. Each operand anchors and replays on its own instead, and the join runs at every + // step over those per-series results. + let plan = build_at_modifier_plan("some_metric @ 300 + some_metric @ 0", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // One anchoring node and one replay per operand... + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 4, + "{plan_str}" + ); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 2, + "{plan_str}" + ); + // ... and the two operands are anchored at different timestamps, so both select their own + // sample. + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..0], lookback=[1000]") + .count(), + 2, + "{plan_str}" + ); +} + +/// Only a direct call over an anchored range selector is promoted. A value function or a unary +/// operator over an anchored *instant* selector needs no promotion: the selector anchors and +/// replays its sample per series on its own, and the operator above it is row-wise, so it can be +/// evaluated at every step over that replay. +#[tokio::test] +async fn at_modifier_does_not_promote_value_calls_over_anchored_selectors() { + for (query, value_expr) in [ + ("abs(some_metric @ 300)", "abs(some_metric.field_0)"), + ("-some_metric @ 300", "(- some_metric.field_0)"), + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The scan is limited to the anchored sample (the lookback delta of this test is 1s). + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(299001, None) AND some_metric.timestamp <= TimestampMillisecond(300000, None)" + ), + "{query}:\n{plan_str}" + ); + // ... and the result is replayed at every step by the selector itself: the anchored + // selection, then the grid replay, with no promoted subtree on top of the operator. + assert!( + !plan_str.starts_with("PromInstantManipulate"), + "{query}:\n{plan_str}" + ); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..0], lookback=[1000]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .unwrap_or_else(|| panic!("replay node:\n{plan_str}")); + let value = plan_str + .find(value_expr) + .unwrap_or_else(|| panic!("no `{value_expr}` projection in:\n{plan_str}")); + assert!( + value < replay, + "`{value_expr}` must be evaluated above the per-series replay:\n{plan_str}" + ); + } +} + +/// A call whose argument merges series (an aggregation or a join) or whose own output reorders +/// the whole vector (`sort*`, the histogram folds) is never the promoted root either: it is +/// planned as usual over the leaf-level anchoring of its selectors, which replays every selector +/// per series and keeps the one-series-per-batch layout the replay needs. +#[tokio::test] +async fn at_modifier_keeps_multi_series_roots_out_of_promoted_subtree() { + // An aggregation below the call emits one row per group in shared batches, so the call is + // not promoted: the replay of the anchored selector stays below the aggregate. + let plan = build_at_modifier_plan("abs(sum(some_metric @ 300))", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + let aggregate = plan_str.find("Aggregate:").expect("aggregate node"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + aggregate < replay, + "the aggregation must stay above the replay:\n{plan_str}" + ); + + // A join of two anchored selectors below the call: neither the join nor the call is + // promoted, so each operand is replayed on its own. + let plan = build_at_modifier_plan("abs(some_metric @ 300 + some_metric @ 0)", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.starts_with("PromInstantManipulate"), "{plan_str}"); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 2, + "{plan_str}" + ); + + // A call that reorders the whole vector keeps its sort above the per-series replay. + let plan = build_at_modifier_plan("sort_by_label(some_metric @ 300, \"tag_0\")", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.starts_with("Sort:"), "{plan_str}"); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 1, + "{plan_str}" + ); +} + +/// `label_join` rewrites the labels of its input series, so it is never the promoted root: the +/// anchored selector keeps replaying one series per batch, and the join runs at every step above +/// that replay (see [`Self::promotes_anchored_range_call`]). Promoting it would replay the +/// joined rows through the labels the join just rewrote, which merges the distinct input series +/// into one timeline. +#[tokio::test] +async fn at_modifier_does_not_promote_label_join() { + for query in [ + // Directly above the anchored instant selector... + "label_join(some_metric @ 300, \"tag_0\", \"-\", \"tag_0\", \"tag_0\")", + // ... and below another call, which is planned as usual over the join. + "abs(label_join(some_metric @ 300, \"tag_0\", \"-\", \"tag_0\", \"tag_0\"))", + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The join is not wrapped in a replay of its own; the only replay over the grid is the + // one of the anchored selector... + assert!( + !plan_str.starts_with("PromInstantManipulate"), + "{query}:\n{plan_str}" + ); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + // ... and the projected join stays above it, evaluated at every step. + let join = plan_str + .find("concat_ws(") + .unwrap_or_else(|| panic!("no `label_join` projection in:\n{plan_str}")); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + join < replay, + "`label_join` must be evaluated above the per-series replay:\n{plan_str}" + ); + } + + // A range call below the join is still promoted on its own: the anchored window is folded + // once per series, and the join above it is evaluated at every step over that replay. + let plan = build_at_modifier_plan( + "label_join(rate(some_metric[5m] @ 300), \"tag_0\", \"-\", \"tag_0\", \"tag_0\")", + 0, + 1000, + ) + .await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + assert!(!plan_str.starts_with("PromInstantManipulate"), "{plan_str}"); + let join = plan_str.find("concat_ws(").expect("join projection"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + join < replay, + "the join must be evaluated above the replay of the range call:\n{plan_str}" + ); +} + +#[test] +fn at_modifier_rejects_subtraction_overflow() { + for (anchor, offset) in [(i64::MAX, -1), (i64::MIN, 1)] { + let err = PromPlanner::anchor_sub(anchor, offset).unwrap_err(); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + assert!( + err.to_string() + .contains("Timestamp out of range for the `@` modifier"), + "{err}" + ); + } + assert_eq!(PromPlanner::anchor_sub(-1, 1).unwrap(), -2); +} + +/// `@` beyond the representable millisecond range is rejected instead of silently wrapping. +/// +/// `@ 1e16` is 10^19 milliseconds, beyond `i64::MAX`. A Unix `SystemTime` can hold it, so +/// the planner rejects the anchor it cannot represent. A Windows `SystemTime` tops out +/// below `i64::MAX` milliseconds, so the same literal is already rejected while parsing. +#[tokio::test] +async fn at_modifier_rejects_unrepresentable_timestamp() { + #[cfg(windows)] + { + let err = parser::parse("some_metric @ 1e16").unwrap_err(); + assert!( + err.to_string() + .contains("timestamp out of bounds for @ modifier"), + "{err}" + ); + } + + #[cfg(not(windows))] + { + let eval_stmt = build_eval_stmt("some_metric @ 1e16"); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let err = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Timestamp out of range for the `@` modifier"), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + } +} + +#[tokio::test] +async fn default_binary_join_uses_tsid_when_available() { + let eval_stmt = build_eval_stmt("some_metric / some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!( + !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn reject_binary_fill_modifiers() { + let state = build_query_engine_state(); + + for query in [ + "some_metric + fill(0) some_alt_metric", + "some_metric + fill_left(0) some_alt_metric", + "some_metric + fill_right(0) some_alt_metric", + "(some_metric + fill(0) some_alt_metric) + some_metric", + ] { + let eval_stmt = build_eval_stmt(query); + let table_provider = build_test_table_provider(&[], 0, 0).await; + let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) + .await + .unwrap_err(); + + assert!( + matches!( + &err, + crate::promql::error::Error::UnsupportedExpr { name, .. } + if name == "PromQL fill modifiers" + ), + "{err}" + ); + } +} + +#[tokio::test] +async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { + for query in [ + "timestamp(some_metric) / some_metric", + "some_metric / timestamp(some_metric)", + ] { + let eval_stmt = build_eval_stmt(query); + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}"); + assert!( + plan_str.contains("lhs.tag_0 = rhs.tag_0"), + "{query}: {plan_str}" + ); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), + "{query}: {plan_str}" + ); + } +} + +#[tokio::test] +async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() { + let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job"); + + let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ + ( + (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()), + 2, + 1, + ), + ( + (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()), + 1, + 1, + ), + ]) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + + assert!( + plan_str.contains("Boolean(false)") || plan_str.contains("false"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn tsid_is_preserved_for_nested_default_binary_joins() { + let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_third_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn repeated_tsid_binary_operand_reuses_leaf_plan() { + let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn repeated_tsid_binary_operand_reuses_shorter_field_side() { + let eval_stmt = + build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let value_columns = field_names + .iter() + .filter(|name| { + *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME + }) + .count(); + assert_eq!(value_columns, 1, "{field_names:?}"); + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1025)") + .count(), + 1, + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn binary_island_reuses_self_operand_without_join() { + let eval_stmt = build_eval_stmt("some_metric / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_reuses_leaf_across_two_branches() { + let eval_stmt = + build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_generated_alias_avoids_user_column_names() { + let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_fields( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + &["prom_v0", "__prom_v0"], + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan.schema().field_names(); + assert!(field_names.iter().any(|name| name.ends_with(".prom_v0"))); + assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0"))); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_clears_qualifier_for_nested_unary_projection() { + let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_keeps_distinct_matcher_leaves() { + let eval_stmt = build_eval_stmt( + "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}", + ); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_keeps_offset_leaves_distinct() { + let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_falls_back_for_group_modifier() { + let eval_stmt = + build_eval_stmt("(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_falls_back_for_comparison_filter() { + let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn tsid_binary_join_uses_shorter_field_side() { + let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let value_columns = field_names + .iter() + .filter(|name| { + *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME + }) + .count(); + assert_eq!(value_columns, 1, "{field_names:?}"); +} + +#[tokio::test] +async fn comparison_binary_join_uses_shorter_field_side() { + let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + assert!( + field_names.iter().any(|name| name == "field_0"), + "{field_names:?}" + ); + assert!( + !field_names.iter().any(|name| name == "field_1"), + "{field_names:?}" + ); +} + +#[tokio::test] +async fn label_matching_modifier_disables_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn ignoring_absent_label_keeps_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() { + let eval_stmt = build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_full_label_set_keeps_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_partial_label_set_disables_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ + ( + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + 2, + 1, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + 1, + 1, + ), + ]) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { + let eval_stmt = build_eval_stmt("some_metric > some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::from_eval_stmt(&eval_stmt).unwrap(), + promql_annotations: None, + }; + let plan = planner + .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), + "{plan_str}" + ); + assert!(planner.ctx.use_tsid, "{plan_str}"); +} + +#[tokio::test] +async fn comparison_bool_binary_join_uses_tsid_when_available() { + let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn scalar_count_count_range_keeps_full_window() { + let plan_str = build_optimized_tsid_plan( + "scalar(count(count(some_metric) by (tag_0)))", + 1, + 1, + 100_000, + 1, + ) + .await; + assert!(plan_str.contains("ScalarCalculate: tags=[]")); + assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]")); + assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]")); +} + +#[tokio::test] +async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() { + let plan_str = build_optimized_tsid_plan( + "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))", + 2, + 1, + 10, + 300, + ) + .await; + assert!(plan_str.contains("Distinct:"), "{plan_str}"); +} + +#[tokio::test] +async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() { + assert_nested_count_rewrite_applies( + "count(count(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]" + ) + .await; +} + +#[tokio::test] +async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() { + assert_nested_count_rewrite_applies( + "count(sum(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]" + ) + .await; +} + +#[tokio::test] +async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() { + for (query, expected_outer_agg) in [ + ( + "count(avg(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]", + ), + ( + "count(min(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]", + ), + ( + "count(max(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]", + ), + ( + "count(stddev(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]", + ), + ( + "count(stdvar(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]", + ), + ] { + assert_nested_count_rewrite_applies(query, expected_outer_agg).await; + } +} + +#[tokio::test] +async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() { + let count_plan = + build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1).await; + assert!( + !count_plan.contains("some_metric.field_0 IS NOT NULL"), + "{count_plan}" + ); + + for query in [ + "count(sum(some_metric) by (tag_0))", + "count(avg(some_metric) by (tag_0))", + "count(min(some_metric) by (tag_0))", + "count(max(some_metric) by (tag_0))", + "count(stddev(some_metric) by (tag_0))", + "count(stdvar(some_metric) by (tag_0))", + ] { + let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; + assert!( + plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"), + "{query}: {plan_str}" + ); + } +} + +#[tokio::test] +async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() { + assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await; + assert_nested_count_rewrite_missing("count(sum(irate(some_metric[1h])) by (tag_0))", 2, 300) + .await; +} + +#[tokio::test] +async fn physical_table_name_is_not_leaked_in_plan() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("TableScan: phy"), "{plan}"); + assert!(plan_str.contains("SubqueryAlias: some_metric")); + assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)")); + assert!(!plan_str.contains("TableScan: some_metric")); +} + +#[tokio::test] +async fn sum_without_does_not_group_by_tsid() { + let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn topk_without_does_not_partition_by_tsid() { + let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let window_line = plan_str + .lines() + .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) + .unwrap(); + let partition_by = window_line + .split("PARTITION BY [") + .nth(1) + .and_then(|s| s.split("] ORDER BY").next()) + .unwrap(); + assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn sum_by_does_not_group_by_tsid() { + let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn aggregate_over_binary_time_function_expr() { + for op in ["sum", "min", "max", "avg"] { + let prom_expr = parser::parse(&format!( + "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)" + )) + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 1, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(aggr_line.contains(op), "{plan_str}"); + assert!(aggr_line.contains("first_value"), "{plan_str}"); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME }) + ); + } +} + +#[tokio::test] +async fn topk_by_does_not_partition_by_tsid() { + let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let window_line = plan_str + .lines() + .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) + .unwrap(); + let partition_by = window_line + .split("PARTITION BY [") + .nth(1) + .and_then(|s| s.split("] ORDER BY").next()) + .unwrap(); + assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn selector_matcher_on_tsid_does_not_use_internal_column() { + let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet) { + if let LogicalPlan::Filter(filter) = plan { + datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap(); + } + for input in plan.inputs() { + collect_filter_cols(input, out); + } + } + + let mut filter_cols = HashSet::new(); + collect_filter_cols(&plan, &mut filter_cols); + assert!( + !filter_cols + .iter() + .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME) + ); +} + +#[tokio::test] +async fn tsid_is_not_used_when_physical_table_is_missing() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let catalog_list = MemoryCatalogManager::with_default_setup(); + + // Register a metric engine logical table referencing a missing physical table. + let mut columns = vec![ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + "field_0".to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let schema = Arc::new(Schema::new(columns)); + let mut options = table::requests::TableOptions::default(); + options + .extra_options + .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string()); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .engine(METRIC_ENGINE_NAME.to_string()) + .options(options) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(1024) + .name("some_metric") + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "some_metric".to_string(), + table_id: 1024, + table, + }) + .unwrap(); + + let table_provider = DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); + assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); +} + +#[tokio::test] +async fn tsid_is_carried_only_when_aggregate_preserves_label_set() { + let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("first_value") && plan_str.contains("__tsid")); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); + + // Merging aggregate: label set is reduced, tsid should not be carried. + let prom_expr = parser::parse("sum(some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("first_value")); +} + +#[tokio::test] +async fn or_operator_with_unknown_metric_does_not_require_tsid() { + let prom_expr = parser::parse("unknown_metric or some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); +} + +#[tokio::test] +async fn aggregate_avg() { + do_aggregate_expr_plan("avg", "avg").await; +} + +#[tokio::test] +#[should_panic] // output type doesn't match +async fn aggregate_count() { + do_aggregate_expr_plan("count", "count").await; +} + +#[tokio::test] +async fn aggregate_min() { + do_aggregate_expr_plan("min", "min").await; +} + +#[tokio::test] +async fn aggregate_max() { + do_aggregate_expr_plan("max", "max").await; +} + +#[tokio::test] +async fn aggregate_group() { + // Regression test for `group()` aggregator. + // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"})) + // should be plannable, and `group()` should produce constant 1 for each group. + let prom_expr = parser::parse( + "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))", + ) + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "kubernetes_build_info".to_string(), + )], + &["cluster", "service", "job"], + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("max(Float64(1")); +} + +#[tokio::test] +async fn aggregate_stddev() { + do_aggregate_expr_plan("stddev", "stddev_pop").await; +} + +#[tokio::test] +async fn aggregate_stdvar() { + do_aggregate_expr_plan("stdvar", "var_pop").await; +} + +// TODO(ruihang): add range fn tests once exprs are ready. + +// { +// input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}", +// expected: &BinaryExpr{ +// Op: ADD, +// LHS: &VectorSelector{ +// Name: "a", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// RHS: &VectorSelector{ +// Name: "sum", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// VectorMatching: &VectorMatching{}, +// }, +// }, +#[tokio::test] +async fn binary_op_column_column() { + let prom_expr = + parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\ + \n Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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(\"foo\") AND 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]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +async fn indie_query_plan(query: &str) -> String { + let prom_expr = parser::parse(query).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + "greptime_private".to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + plan.display_indent_schema().to_string() +} + +async fn indie_query_plan_compare>(query: &str, expected: T) { + let plan = indie_query_plan(query).await; + assert_eq!(plan, expected.as_ref()); +} + +#[tokio::test] +async fn binary_op_literal_column() { + let query = r#"1 + some_metric{tag_0="bar"}"#; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn binary_op_literal_literal() { + let query = r#"1 + 1"#; + let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N] + TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn simple_bool_grammar() { + let query = "some_metric != bool 1.2345"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn bool_with_additional_arithmetic() { + let query = "some_metric + (1 == bool 2)"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn simple_unary() { + let query = "-some_metric"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn increase_aggr() { + let query = "increase(some_metric[5m])"; + let expected = String::from( + "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn predict_linear_injects_the_eval_timestamp() { + // The fourth argument is the evaluation instant of each row, which the fold offset + // recovers from the row's time index (here: no `@` and no `offset`, so the step itself). + let query = "predict_linear(some_metric[5m], 60)"; + let expected = String::from( + "Filter: prom_predict_linear(timestamp_range,field_0,Float64(60)) IS NOT NULL [timestamp:Timestamp(ms), prom_predict_linear(timestamp_range,field_0,Float64(60)):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_predict_linear(timestamp_range, field_0, CAST(Float64(60) AS Int64), CAST(CAST(some_metric.timestamp AS Int64) + Int64(0) AS Timestamp(ms))) AS prom_predict_linear(timestamp_range,field_0,Float64(60)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_predict_linear(timestamp_range,field_0,Float64(60)):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +/// The evaluation instant follows the selector's fold offset, not the projection's: an +/// `@`-anchored selector is folded with `at_offset`, which is what the window's timestamps were +/// shifted by. The evaluation starts at 0s here, so `@ 100` anchors 100s in the future and +/// yields `at_offset` = -100000ms. +#[tokio::test] +async fn predict_linear_eval_ts_follows_the_fold_offset() { + for (query, expected_offset) in [ + ( + "predict_linear(some_metric[5m] @ 100, 60)", + "Int64(-100000)", + ), + ( + "predict_linear(some_metric[5m] offset 2m, 60)", + "Int64(120000)", + ), + ] { + let plan = indie_query_plan(query).await; + assert!( + plan.contains(&format!( + "some_metric.timestamp AS Int64) + {expected_offset}" + )), + "{query}\n{plan}" + ); + } +} + +async fn native_histogram_plan(query: &str) -> String { + let table_provider = build_test_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent_schema().to_string() +} + +#[tokio::test] +async fn native_histogram_count_uses_native_udf() { + let plan = native_histogram_plan("histogram_count(some_metric)").await; + + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); +} + +#[tokio::test] +async fn timestamp_filters_native_histogram_stale_marker_before_projection() { + let mut stale = direct_or_histogram(); + stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + let table = operator_metric_table( + "stale_histogram", + 2_100, + "a", + None, + DirectOrValue::NativeHistogram(stale), + ); + let catalog = MemoryCatalogManager::with_default_setup(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "stale_histogram".to_string(), + table_id: 2_100, + table, + }) + .unwrap(); + let provider = DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + provider, + &operator_eval_stmt("timestamp(stale_histogram)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); +} + +#[tokio::test] +async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { + let histograms = build_histogram_array(&[None]); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new( + greptime_native_histogram(), + histograms.data_type().clone(), + true, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + histograms, + Arc::new(Float64Array::from(vec![f64::from_bits( + PROMETHEUS_STALE_NAN_BITS, + )])), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + Vec::new(), + "timestamp".to_string(), + input, + )), + }); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + 1_000, + 1_000, + 5_000, + 1_000, + 0, + "timestamp".to_string(), + Vec::new(), + Some(greptime_native_histogram().to_string()), + input, + )), + }); + // Match timestamp()'s parent projection, which otherwise prunes the companion lane. + let plan = LogicalPlanBuilder::from(input) + .project([col("timestamp")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); +} + +#[tokio::test] +async fn native_histogram_rate_can_feed_count() { + let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; + + assert!(plan.contains("prom_native_histogram_rate"), "{plan}"); + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); +} + +#[tokio::test] +async fn native_histogram_quantile_skips_classic_fold() { + let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); + assert!(plan.contains("some_metric.le"), "{plan}"); + // The phi literal is threaded into the native quantile UDF as its second argument. + assert!(plan.contains("Float64(0.9)"), "{plan}"); + // The empty-values filter drops NULL quantile results so the output is empty + // when all native histogram samples are dropped. + assert!(plan.contains("IS NOT NULL"), "{plan}"); +} + +#[tokio::test] +async fn mixed_native_histogram_quantile_uses_histogram_field() { + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_quantile(0.9, some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + + assert!( + plan.contains("prom_native_histogram_quantile(greptime_native_histogram"), + "{plan}" + ); + assert!(!plan.contains("EmptyRelation"), "{plan}"); +} + +#[tokio::test] +async fn mixed_histogram_helpers_execute_classic_and_native_samples() { + let state = build_query_engine_state(); + for (query, expected) in [ + ( + "histogram_quantile(0.5, mixed_histogram)", + vec![("classic", 1.0), ("native", 0.0)], + ), + ( + "histogram_fraction(-Inf, +Inf, mixed_histogram)", + vec![("classic", 1.0), ("native", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); + assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let tags = batch + .column_by_name("service.name") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) + }) + .collect::>(); + actual.sort_by_key(|(tag, _)| *tag); + assert_eq!(actual, expected, "{query}"); + } +} + +#[tokio::test] +async fn mixed_histogram_helpers_report_annotations() { + let state = build_query_engine_state(); + let mut native_histogram = direct_or_histogram(); + native_histogram.count = 2.0; + native_histogram.sum = f64::NAN; + for (native_tag, expected_rows, expected_warnings, expected_infos) in [ + ( + "classic", + 0, + vec!["vector contains a mix of classic and native histograms"], + vec![], + ), + ( + "native", + 2, + vec![], + vec!["input to histogram_quantile has NaN observations, result is skewed higher"], + ), + ] { + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + classic_and_native_histogram_table_provider(native_tag, None, native_histogram.clone()), + &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows + ); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!(warnings, expected_warnings); + assert_eq!(infos, expected_infos); + } +} + +#[tokio::test] +async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { + let state = build_query_engine_state(); + let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); + stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "classic", + Some("native"), + direct_or_histogram(), + ), + &stmt, + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert_eq!( + plan_text.matches("TableScan: mixed_histogram").count(), + 1, + "{plan_text}" + ); + + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let le = batch + .column_by_name(LE_COLUMN_NAME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = batch + .column_by_name("timestamp") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| { + ( + timestamps.value(row), + (!le.is_null(row)).then(|| le.value(row).to_string()), + values.value(row), + ) + }) + }) + .collect::>(); + actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); + assert_eq!( + actual, + vec![ + (1_000, None, 1.0), + (1_000, Some("native".to_string()), 0.0), + (2_000, None, 1.0), + (2_000, Some("native".to_string()), 0.0), + ] + ); +} + +#[tokio::test] +async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { + let state = build_query_engine_state(); + for native_le in [None, Some("native")] { + for query in [ + "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", + "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "native", + native_le, + direct_or_histogram(), + ), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "native_le={native_le:?}, query={query}" + ); + } + } +} + +#[tokio::test] +async fn native_histogram_quantile_rejects_multi_field_input() { + let table_provider = build_test_multi_histogram_table_provider("some_metric").await; + let result = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_quantile(0.9, some_metric)"), + &build_query_engine_state(), + ) + .await; + + let err = result.expect_err("histogram_quantile on two native histogram fields must fail"); + assert!( + err.to_string() + .contains("Multi fields calculation is not supported in histogram_quantile"), + "{err}" + ); +} + +#[tokio::test] +async fn native_histogram_topk_uses_drop_udf() { + let plan = native_histogram_plan("topk(1, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}"); + assert!( + plan.contains("Filter: prom_native_histogram_drop_float") && plan.contains("IS NOT NULL"), + "{plan}" + ); +} + +#[tokio::test] +async fn mixed_or_topk_bottomk_ignore_native_histograms() { + for op in ["topk", "bottomk"] { + let collector = PromqlAnnotationCollector::default(); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), + "{plan:?}" + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!( + infos, + vec![format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + )] + ); + } +} + +#[tokio::test] +async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { + let plan = native_histogram_plan("scalar(some_metric)").await; + + assert!(plan.contains("ScalarCalculate"), "{plan}"); + assert!(plan.contains("Filter: Boolean(false)"), "{plan}"); + assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); +} + +#[tokio::test] +async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() { + for function in ["sort", "sort_desc"] { + let plan = native_histogram_plan(&format!("{function}(some_metric)")).await; + + assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); + assert!( + !plan.contains(&format!("Sort: {}", greptime_native_histogram())), + "{plan}" + ); + assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); + } + + for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] { + let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await; + + assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}"); + assert!(plan.contains(greptime_native_histogram()), "{plan}"); + assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); + } +} + +#[tokio::test] +async fn unsupported_native_histogram_functions_use_drop_udf() { + for query in [ + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "round(some_metric)", + "abs(some_metric)", + ] { + let plan = native_histogram_plan(query).await; + + assert!( + plan.contains("prom_native_histogram_drop_float"), + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn native_histogram_absent_over_time_uses_native_udf() { + let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await; + + assert!( + plan.contains("prom_native_histogram_absent_over_time"), + "{plan}" + ); +} + +#[tokio::test] +async fn native_histogram_all_function_arms_route_correctly() { + // Every native-histogram match arm in `create_function_expr` must route to the + // expected UDF when all field columns are native histograms. `holt_winters` shares + // the `double_exponential_smoothing` arm but is not registered in the promql + // parser (0.10), so it cannot be exercised through a query string. + let cases = [ + // Range functions routed to native histogram UDFs. + ( + "increase(some_metric[5m])", + "prom_native_histogram_increase", + ), + ("rate(some_metric[5m])", "prom_native_histogram_rate"), + ("delta(some_metric[5m])", "prom_native_histogram_delta"), + ("idelta(some_metric[5m])", "prom_native_histogram_idelta"), + ("irate(some_metric[5m])", "prom_native_histogram_irate"), + ("resets(some_metric[5m])", "prom_native_histogram_resets"), + ("changes(some_metric[5m])", "prom_native_histogram_changes"), + ( + "avg_over_time(some_metric[5m])", + "prom_native_histogram_avg_over_time", + ), + ( + "sum_over_time(some_metric[5m])", + "prom_native_histogram_sum_over_time", + ), + ( + "count_over_time(some_metric[5m])", + "prom_native_histogram_count_over_time", + ), + ( + "last_over_time(some_metric[5m])", + "prom_native_histogram_last_over_time", + ), + ( + "present_over_time(some_metric[5m])", + "prom_native_histogram_present_over_time", + ), + // Unsupported functions dropped with the float-null UDF. + ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"), + ( + "min_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "max_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "stddev_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "stdvar_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "quantile_over_time(0.9, some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "predict_linear(some_metric[5m], 60)", + "prom_native_histogram_drop_float", + ), + ( + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + "prom_native_histogram_drop_float", + ), + ("round(some_metric)", "prom_native_histogram_drop_float"), + ("rad(some_metric)", "prom_native_histogram_drop_float"), + ("deg(some_metric)", "prom_native_histogram_drop_float"), + ("sgn(some_metric)", "prom_native_histogram_drop_float"), + // Instant helper functions routed to native histogram UDFs. + ( + "histogram_count(some_metric)", + "prom_native_histogram_count", + ), + ("histogram_sum(some_metric)", "prom_native_histogram_sum"), + ("histogram_avg(some_metric)", "prom_native_histogram_avg"), + ( + "histogram_stddev(some_metric)", + "prom_native_histogram_stddev", + ), + ( + "histogram_stdvar(some_metric)", + "prom_native_histogram_stdvar", + ), + ( + "histogram_fraction(-2 + 1, 2 / 2, some_metric)", + "prom_native_histogram_fraction", + ), + ]; + + for (query, expected_udf) in cases { + let plan = native_histogram_plan(query).await; + assert!(plan.contains(expected_udf), "{query}\n{plan}"); + if query.starts_with("histogram_fraction") { + assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); + } + } +} + +#[tokio::test] +async fn mixed_native_histogram_ranges_use_coordinated_udfs() { + let dual_output = [ + "increase(some_metric[5m])", + "rate(some_metric[5m])", + "delta(some_metric[5m])", + "idelta(some_metric[5m])", + "irate(some_metric[5m])", + "avg_over_time(some_metric[5m])", + "sum_over_time(some_metric[5m])", + "last_over_time(some_metric[5m])", + ]; + let float_output = [ + "resets(some_metric[5m])", + "changes(some_metric[5m])", + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "max_over_time(some_metric[5m])", + "count_over_time(some_metric[5m])", + "absent_over_time(some_metric[5m])", + "present_over_time(some_metric[5m])", + "stddev_over_time(some_metric[5m])", + "stdvar_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + ]; + + for query in dual_output.iter().chain(float_output.iter()) { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); + assert_eq!( + plan.contains("prom_mixed_range_histogram"), + dual_output.contains(query), + "{query}\n{plan}" + ); + } + + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; + assert_eq!(plan, expected); +} + +#[tokio::test] +async fn mixed_native_histogram_predict_linear_forwards_the_eval_timestamp() { + let query = "predict_linear(some_metric[5m], 60)"; + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + + assert!( + plan.contains( + "prom_mixed_range_float(Utf8(\"predict_linear\"), timestamp_range, greptime_value, greptime_native_histogram, CAST(Float64(60) AS Int64), CAST(CAST(some_metric.timestamp AS Int64) + Int64(0) AS Timestamp(ms)))" + ), + "{query}\n{plan}" + ); +} + +#[tokio::test] +async fn mixed_native_histogram_rate_executes_real_ranges() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + Field::new( + greptime_native_histogram(), + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), + Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), + build_histogram_array(&[None, Some(direct_or_histogram()), None]), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let collector = PromqlAnnotationCollector::default(); + let mut planner = PromPlanner { + table_provider: build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await, + ctx: PromPlannerContext { + start: 3000, + end: 3000, + interval: 1000, + range: Some(3000), + time_index_column: Some("timestamp".to_string()), + field_columns: vec![ + greptime_native_histogram().to_string(), + greptime_value().to_string(), + ], + ..Default::default() + }, + promql_annotations: Some(collector.clone()), + }; + let input = LogicalPlan::Extension(Extension { + node: Arc::new( + RangeManipulate::new( + 3000, + 3000, + 1000, + 0, + 3000, + "timestamp".to_string(), + planner.ctx.field_columns.clone(), + input, + ) + .unwrap(), + ), + }); + let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { + unreachable!() + }; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state, None) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = Vec::new(); + collector.append_to(&mut warnings, &mut Vec::new()); + assert!( + warnings + .iter() + .any(|warning| warning.contains("mix of float and native histogram")) + ); +} + +#[tokio::test] +async fn native_histogram_mixed_field_table_behaves() { + // Exercise function planning after float and histogram samples have already been + // represented as alternative nullable fields. Histogram functions must select the + // histogram field without adding a NULL float field that would reject every row. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_count(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("prom_native_histogram_count"), + "{plan_str}" + ); + assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}"); + assert!( + plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"), + "{plan_str}" + ); + + // Value sorting keeps the float column and never sorts by the histogram column. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("sort(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("greptime_value ASC NULLS FIRST"), + "{plan_str}" + ); + assert!( + !plan_str.contains("greptime_native_histogram ASC"), + "{plan_str}" + ); + + // scalar() ignores histogram samples and evaluates only the float field. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("scalar(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("ScalarCalculate"), "{plan_str}"); + assert!( + plan_str.contains("greptime_value IS NOT NULL"), + "{plan_str}" + ); + + // Functions that preserve both alternative fields keep rows with either sample type. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + let filter = plan_str.lines().next().unwrap(); + assert!( + filter.starts_with("Filter: ") + && filter.contains("greptime_native_histogram IS NOT NULL") + && filter.contains(" OR ") + && filter.contains("greptime_value IS NOT NULL"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn less_filter_on_value() { + let query = "some_metric < 1.2345"; + let expected = String::from( + "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn count_over_time() { + let query = "count_over_time(some_metric[5m])"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +/// The outer `PromRangeManipulate` from a subquery must be preceded by +/// `Sort` + `PromSeriesDivide`. +#[tokio::test] +async fn count_over_time_subquery() { + let query = "count_over_time(some_metric[10m:1m])"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \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 PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + indie_query_plan_compare(query, expected).await; +} + +/// `offset` on a subquery must shift the inner evaluation window back and be +/// carried into the outer range manipulation. See +/// . +#[tokio::test] +async fn count_over_time_subquery_with_offset() { + let query = "count_over_time(some_metric[10m:1m] offset 5m)"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[300000], time index=[timestamp], filter NaN: [false] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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 PromInstantManipulate: range=[-840000..99700000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-840999, None) AND some_metric.timestamp <= TimestampMillisecond(99700000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + indie_query_plan_compare(query, expected).await; +} + +/// Time-varying samples on and off the subquery step grid expose incorrect window folding. +async fn build_subquery_offset_probe_provider() -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let probe_rows = (0..=6_i64) + .map(|step| (step * 10_000, (step + 1) as f64)) + .collect::>(); + let offgrid_rows = (0..6_i64) + .map(|step| (step * 10_000 + 5_000, (step + 1) as f64)) + .collect::>(); + + for (table_id, name, rows) in [ + (7_001_u32, "subquery_offset_probe", &probe_rows), + (7_002_u32, "subquery_offset_offgrid", &offgrid_rows), + ] { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "host".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "val".to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec!["a"; rows.len()])), + Arc::new(TimestampMillisecondArray::from_iter_values( + rows.iter().map(|(ts, _)| *ts), + )), + Arc::new(Float64Array::from_iter_values( + rows.iter().map(|(_, value)| *value), + )), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + name, + GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), + table_id, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(3) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(table_id) + .name(name) + .meta(table_meta) + .build() + .unwrap(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: name.to_string(), + table_id, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +/// Plans on a seconds-based evaluation grid. +async fn plan_subquery_offset_probe( + query: &str, + start: u64, + end: u64, + step: u64, +) -> Result { + let eval_stmt = EvalStmt { + expr: parser::parse(query).unwrap(), + start: UNIX_EPOCH + Duration::from_secs(start), + end: UNIX_EPOCH + Duration::from_secs(end), + interval: Duration::from_secs(step), + lookback_delta: Duration::from_secs(300), + }; + PromPlanner::stmt_to_plan( + build_subquery_offset_probe_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await +} + +/// Evaluates seconds-based bounds and returns sorted `(timestamp_ms, value)` rows. +async fn run_subquery_offset_probe( + query: &str, + start: u64, + end: u64, + step: u64, +) -> Vec<(i64, f64)> { + let state = build_query_engine_state(); + let plan = plan_subquery_offset_probe(query, start, end, step) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + let mut rows = Vec::new(); + for batch in &batches { + let timestamp_index = batch + .schema() + .fields() + .iter() + .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) + .expect("no timestamp column"); + let value_index = batch + .schema() + .fields() + .iter() + .position(|field| field.data_type() == &ArrowDataType::Float64) + .expect("no Float64 value column"); + let timestamps = batch + .column(timestamp_index) + .as_any() + .downcast_ref::() + .expect("timestamp column is not a millisecond timestamp"); + let values = batch + .column(value_index) + .as_any() + .downcast_ref::() + .expect("value column is not Float64"); + rows.extend( + timestamps + .iter() + .zip(values.iter()) + .map(|(timestamp, value)| (timestamp.unwrap(), value.unwrap())), + ); + } + rows.sort_by_key(|(timestamp, _)| *timestamp); + rows +} + +#[tokio::test] +async fn subquery_own_offset_shifts_the_inner_window_but_keeps_evaluation_timestamps() { + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] offset 30s)", + 45, + 45, + 1 + ) + .await, + vec![(45_000, 3.0)], + ); + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_probe[20s:10s])", 15, 15, 1).await, + vec![(15_000, 3.0)], + ); + + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] offset 30s)", + 30, + 60, + 30 + ) + .await, + vec![(30_000, 1.0), (60_000, 7.0)], + ); + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_probe[20s:10s])", 0, 30, 30).await, + vec![(0, 1.0), (30_000, 7.0)], + ); + + // Negative offsets look ahead while retaining outer timestamps. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] offset -30s)", + 0, + 30, + 30 + ) + .await, + vec![(0, 7.0), (30_000, 13.0)], + ); + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_probe[20s:10s])", 30, 60, 30) + .await, + vec![(30_000, 7.0), (60_000, 13.0)], + ); +} + +/// Expectations match Prometheus v3.14.0. +#[tokio::test] +async fn subquery_inner_points_are_anchored_on_absolute_step_multiples() { + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[20s:10s])", 45, 45, 1) + .await, + vec![(45_000, 7.0)], + ); + + // An off-grid offset distinguishes absolute anchoring from re-phasing. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s)", + 30, + 30, + 1 + ) + .await, + vec![(30_000, 3.0)], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s)", + 45, + 45, + 1 + ) + .await, + vec![(45_000, 7.0)], + ); + + // The window starts before the first sample. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 15s)", + 30, + 30, + 1 + ) + .await, + vec![(30_000, 1.0)], + ); + + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset -5s)", + 15, + 15, + 1 + ) + .await, + vec![(15_000, 3.0)], + ); + + // The negative grid start exercises Euclidean alignment. + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[40s:10s])", 15, 15, 1) + .await, + vec![(15_000, 1.0)], + ); + + // Multi-step evaluation starts off the subquery grid. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s])", + 35, + 65, + 30 + ) + .await, + vec![(35_000, 5.0), (65_000, 11.0)], + ); + + // Child points reach no sample at the first evaluation. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 30s)", + 30, + 60, + 30 + ) + .await, + vec![(60_000, 5.0)], + ); +} + +#[test] +fn subquery_child_window_is_the_child_grid_window() { + let child_window = PromPlanner::subquery_child_window; + // Exclusive lower bound; upper bound is the last parent evaluation step. + assert_eq!( + child_window(30_000, 60_000, 10_000, 0, 20_000, 10_000).unwrap(), + Some((20_000, 60_000)) + ); + assert_eq!( + child_window(45_000, 60_000, 10_000, 0, 20_000, 10_000).unwrap(), + Some((30_000, 55_000)) + ); + assert_eq!( + child_window(25_000, 60_000, 10_000, 0, 20_000, 10_000).unwrap(), + Some((10_000, 55_000)) + ); + assert_eq!( + child_window(35_000, 65_000, 30_000, 0, 20_000, 10_000).unwrap(), + Some((20_000, 65_000)) + ); + assert_eq!( + child_window(0, 55_000, 30_000, 0, 20_000, 10_000).unwrap(), + Some((-10_000, 30_000)) + ); + // A nonempty child window can contribute no parent rows. + assert_eq!( + child_window(35_000, 65_000, 10_000, 0, 5_000, 10_000).unwrap(), + Some((40_000, 65_000)) + ); + // The first child point is past the parent’s last evaluation. + assert_eq!( + child_window(35_000, 44_000, 10_000, 0, 5_000, 10_000).unwrap(), + None + ); + assert_eq!( + child_window(30_000, 30_000, 1_000, 5_000, 20_000, 10_000).unwrap(), + Some((10_000, 25_000)) + ); + assert_eq!( + child_window(15_000, 15_000, 1_000, -5_000, 20_000, 10_000).unwrap(), + Some((10_000, 20_000)) + ); + // Boundaries before the epoch floor towards minus infinity. + assert_eq!( + child_window(-15_000, -15_000, 1_000, 0, 20_000, 10_000).unwrap(), + Some((-30_000, -15_000)) + ); + assert_eq!( + child_window(-45_000, -45_000, 1_000, -10_000, 20_000, 10_000).unwrap(), + Some((-50_000, -35_000)) + ); + // A grid start before the first sample is not clamped to it. + assert_eq!( + child_window(0, 0, 1_000, 0, 20_000, 10_000).unwrap(), + Some((-10_000, 0)) + ); + // A zero parent interval leaves the upper bound unaligned. + assert_eq!( + child_window(0, 55_000, 0, 0, 20_000, 10_000).unwrap(), + Some((-10_000, 55_000)) + ); + + // Only the derived start overflows. + let err = child_window(i64::MIN + 1_000, i64::MIN + 1_000, 1_000, 0, 2_000, 10).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start") && timestamp.contains("-9223372036854776800") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // Only the derived end overflows. + let err = child_window(i64::MAX - 1_000, i64::MAX - 1_000, 1, -2_000, 3_000, 10).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid end") && timestamp.contains("9223372036854776807") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // The wide window is nonempty even though its start overflows. + let err = child_window(i64::MIN, i64::MAX, 10_000, i64::MAX, i64::MAX, 1_000).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start") + ), + "{err}" + ); + // Both derived bounds overflow; report the exact start without clipping. + let err = child_window(i64::MAX - 1_000, i64::MAX - 1_000, 1, i64::MIN, 10, 1).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start") + && timestamp.contains("18446744073709550606") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // Check emptiness before converting overflowing bounds. + assert_eq!( + child_window(i64::MAX, i64::MAX, 1_000, i64::MIN, 0, 1_000).unwrap(), + None + ); +} + +#[tokio::test] +async fn subquery_child_window_without_a_step_reports_no_row() { + // The first child point is after the parent’s only evaluation. + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:10s])", 35, 44, 10) + .await, + vec![], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1)", + 35, + 44, + 10 + ) + .await, + vec![(35_000, 1.0)], + ); + // Multi-step: every step of the grid reports no child point and falls back. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1)", + 35, + 65, + 10 + ) + .await, + vec![(35_000, 1.0), (45_000, 1.0), (55_000, 1.0), (65_000, 1.0)], + ); + // A nonempty child grid can yield no parent rows; a longer range does. + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:10s])", 35, 65, 10) + .await, + vec![], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s])", + 35, + 65, + 10 + ) + .await, + vec![(35_000, 5.0), (45_000, 7.0), (55_000, 9.0), (65_000, 11.0)], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s)", + 35, + 65, + 10 + ) + .await, + vec![(35_000, 5.0), (45_000, 7.0), (55_000, 9.0), (65_000, 11.0)], + ); + // A negative offset makes the short window reach a child point. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[5s:10s] offset -5s)", + 35, + 44, + 10 + ) + .await, + vec![(35_000, 4.0)], + ); + + // An empty outer window suppresses inner results, but not an outer fallback. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s])", + 35, + 44, + 10 + ) + .await, + vec![], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]) or vector(1)", + 35, + 44, + 10 + ) + .await, + vec![(35_000, 1.0)], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1))[5s:10s])", + 35, + 44, + 10 + ) + .await, + vec![], + ); +} + +#[tokio::test] +async fn subquery_rejects_invalid_range_and_zero_step() { + let err = + plan_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:10s])", 44, 35, 10) + .await + .unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::InvalidTimeRange { + start: 44_000, + end: 35_000, + .. + } + ), + "{err}" + ); + let err = plan_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s])", + 44, + 35, + 10, + ) + .await + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + + // `[5s:0]` parses to zero; planning rejects it. + let err = plan_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:0])", 60, 60, 1) + .await + .unwrap_err(); + assert!( + matches!(&err, crate::promql::error::Error::ZeroRangeSelector { .. }), + "{err}" + ); + // An omitted child step inherits the zero parent step. + let err = plan_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:])", 60, 60, 0) + .await + .unwrap_err(); + assert!( + matches!(&err, crate::promql::error::Error::ZeroRangeSelector { .. }), + "{err}" + ); +} + +/// Rejects nonempty child windows whose derived millisecond bounds are unrepresentable. +#[tokio::test] +async fn subquery_rejects_a_child_window_out_of_the_representable_range() { + // Offset is i64::MAX milliseconds; only the derived start overflows. + let err = plan_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[21366776s:10s] offset 292471208y21366775s807ms)", + 0, + 60, + 10, + ) + .await + .unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start = -9223372058221550000") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // Offset is -i64::MAX milliseconds; only the derived end overflows. + let err = plan_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset -292471208y21366775s807ms)", + 0, + 60, + 10, + ) + .await + .unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid end = 9223372036854835807") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); +} + +#[tokio::test] +async fn test_hash_join() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_sum".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["uri", "kubernetes_namespace", "kubernetes_pod_name"], + ) + .await; + // Should be ok + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ + \n Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ + \n SubqueryAlias: http_server_requests_seconds_sum\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_sum\ + \n SubqueryAlias: http_server_requests_seconds_count\ + \n Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_count"; + assert_eq!(plan.to_string(), expected); +} + +#[tokio::test] +async fn test_nested_histogram_quantile() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "greptime_servers_grpc_requests_elapsed_bucket".to_string(), + )], + &["pod", "le", "path", "code", "container"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_histogram_quantile_binary_op() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // Arithmetic applied to a histogram_quantile() result. Regression for #8144: + // HistogramFold used to drop the input column qualifiers, so the binary-op + // projection failed to resolve the qualified tag column. + let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "http_request_duration_seconds_bucket".to_string(), + )], + &["pod", "le"], + ) + .await; + // Should plan without a "No field named ..." error. + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_parse_and_operator() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let cases = [ + r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, + r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, + ]; + + for case in cases { + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "kubelet_volume_stats_used_bytes".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "kubelet_volume_stats_capacity_bytes".to_string(), + ), + ], + &["namespace", "persistentvolumeclaim"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn test_nested_binary_op() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) - + ( + sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m])) + or + vector(0) + )"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "nginx_ingress_controller_requests".to_string(), + )], + &["namespace", "job"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_parse_or_operator() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#" + sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) / + (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100) + or + 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) / + sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#; + + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) + + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 + or + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 + or + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#; + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) + + sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or + (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or + (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "background_waitevent_cnt".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "foreground_waitevent_cnt".to_string(), + ), + ], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#; + let table_provider = build_test_table_provider_with_fields( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "container_cpu_load_average_10s".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "container_spec_cpu_quota".to_string(), + ), + ], + &["cluster_name", "host_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn value_matcher() { + // template + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let cases = [ + // single equal matcher + ( + r#"some_metric{__field__="field_1"}"#, + vec![ + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // two equal matchers + ( + r#"some_metric{__field__="field_1", __field__="field_0"}"#, + vec![ + "some_metric.field_0", + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single not_eq matcher + ( + r#"some_metric{__field__!="field_1"}"#, + vec![ + "some_metric.field_0", + "some_metric.field_2", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // two not_eq matchers + ( + r#"some_metric{__field__!="field_1", __field__!="field_2"}"#, + vec![ + "some_metric.field_0", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // equal and not_eq matchers (no conflict) + ( + r#"some_metric{__field__="field_1", __field__!="field_0"}"#, + vec![ + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // equal and not_eq matchers (conflict) + ( + r#"some_metric{__field__="field_2", __field__!="field_2"}"#, + vec![ + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single regex eq matcher + ( + r#"some_metric{__field__=~"field_1|field_2"}"#, + vec![ + "some_metric.field_1", + "some_metric.field_2", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single regex not_eq matcher + ( + r#"some_metric{__field__!~"field_1|field_2"}"#, + vec![ + "some_metric.field_0", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + ]; + + for case in cases { + let prom_expr = parser::parse(case.0).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let mut fields = plan.schema().field_names(); + let mut expected = case.1.into_iter().map(String::from).collect::>(); + fields.sort(); + expected.sort(); + assert_eq!(fields, expected, "case: {:?}", case.0); + } + + let bad_cases = [ + r#"some_metric{__field__="nonexistent"}"#, + r#"some_metric{__field__!="nonexistent"}"#, + ]; + + for case in bad_cases { + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await; + assert!(plan.is_err(), "case: {:?}", case); + } +} + +#[tokio::test] +async fn custom_schema() { + let query = "some_alt_metric{__schema__=\"greptime_private\"}"; + let expected = String::from( + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; + + let query = "some_alt_metric{__database__=\"greptime_private\"}"; + let expected = String::from( + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; + + let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\ + \n Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \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.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn only_equals_is_supported_for_special_matcher() { + let queries = &[ + "some_alt_metric{__schema__!=\"greptime_private\"}", + "some_alt_metric{__schema__=~\"lalala\"}", + "some_alt_metric{__database__!=\"greptime_private\"}", + "some_alt_metric{__database__=~\"lalala\"}", + ]; + + for query in queries { + let prom_expr = parser::parse(query).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + "greptime_private".to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await; + assert!(plan.is_err(), "query: {:?}", query); + } +} + +#[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")).unwrap(), + 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(); + let columns = vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "field".to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name("metrics".to_string()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "metrics".to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + let plan = PromPlanner::stmt_to_plan( + DfTableSourceProvider::new( + catalog_list.clone(), + false, + QueryContext::arc(), + DummyDecoder::arc(), + true, + ), + &EvalStmt { + expr: parser::parse("metrics{tag = \"1\"}").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }, + &build_query_engine_state(), + ) + .await + .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(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( + catalog_list.clone(), + false, + QueryContext::arc(), + DummyDecoder::arc(), + true, + ), + &EvalStmt { + expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }, + &build_query_engine_state(), + ) + .await + .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(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]" + ); +} + +#[tokio::test] +async fn test_nonexistent_label() { + // template + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"some_metric{nonexistent="hi"}"#; + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_label_join() { + let prom_expr = + parser::parse("label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')") + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = + build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = r#" +Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] + Projection: up.timestamp, up.field_0, nullif(concat_ws(Utf8(","), coalesce(up.tag_1, Utf8("")), coalesce(up.tag_2, Utf8("")), coalesce(up.tag_3, Utf8(""))), Utf8("")) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; + + let ret = plan.display_indent_schema().to_string(); + assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); +} + +#[tokio::test] +async fn test_label_replace() { + let prom_expr = + parser::parse("label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")") + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = + build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = r#" +Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] + Projection: up.timestamp, up.field_0, CASE WHEN regexp_like(coalesce(up.tag_0, Utf8("")), Utf8("^(?s:(.*):.*)$")) THEN nullif(regexp_replace(coalesce(up.tag_0, Utf8("")), Utf8("^(?s:(.*):.*)$"), Utf8("$1")), Utf8("")) ELSE Utf8(NULL) END AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; + + let ret = plan.display_indent_schema().to_string(); + assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); +} + +#[tokio::test] +async fn label_replace_aggregation_queries_plan_successfully() { + let aggregate = r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#; + let queries = [ + aggregate.to_string(), + format!("{aggregate} <= 10"), + format!("{aggregate} * 0.8"), + format!("0.8 * {aggregate}"), + format!("{aggregate} <= {aggregate} * 0.8"), + ]; + let state = build_query_engine_state(); + let mut failures = Vec::new(); + + for query in queries { + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + if let Err(error) = + PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await + { + failures.push(format!("{query}: {error:?}")); + } + } + + assert!(failures.is_empty(), "{}", failures.join("\n")); +} + +#[tokio::test] +async fn test_matchers_to_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = + r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + )], + 3, + 3, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]"; + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_topk_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\ + \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_count_values_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_value_alias() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series"); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = r#" +Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)] + Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; + assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); +} + +#[tokio::test] +async fn test_quantile_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_or_not_exists_table_label() { + let state = build_query_engine_state(); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], + &["job"], + ) + .await; + let raw = PromPlanner::stmt_to_plan( + provider, + &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), + &state, + ) + .await + .unwrap(); + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_0@") + ); + let (optimized, batches) = execute(raw, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); +} + +#[tokio::test] +async fn test_histogram_quantile_missing_le_column() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // Test case: histogram_quantile with a table that doesn't have 'le' column + let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + + // Create a table provider with a table that doesn't have 'le' column + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "non_existent_histogram_bucket".to_string(), + )], + &["pod", "instance"], // Note: no 'le' column + ) + .await; + + // Should return empty result instead of error + let result = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()).await; + + // This should succeed now (returning empty result) instead of failing with "Cannot find column le" + assert!( + result.is_ok(), + "Expected successful plan creation with empty result, but got error: {:?}", + result.err() + ); + + // Verify that the result is an EmptyRelation + let plan = result.unwrap(); + match plan { + LogicalPlan::EmptyRelation(_) => { + // This is what we expect + } + _ => panic!("Expected EmptyRelation, but got: {:?}", plan), + } +} + +#[tokio::test] +async fn test_direct_or_normalizes_missing_match_labels() { + type Case<'a> = ( + Option>, + Option>, + i64, + i64, + &'a [(f64, Option<&'a str>)], + ); + + let modifier = or_modifier("lhs or on(k) rhs"); + #[rustfmt::skip] + let cases: &[Case<'_>] = &[ + (None, None, 1, 1, &[(1.0, None)]), + (None, Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), + (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), + (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), + ]; + for &(left, right, left_ts, right_ts, expected) in cases { + let (optimized, batches) = run( + &matrix_source("lhs", left, left_ts, 1.0), + &matrix_source("rhs", right, right_ts, 2.0), + matrix_context("lhs", left), + matrix_context("rhs", right), + &modifier, + ) + .await; + assert_no_internal_or_keys(optimized.schema()); + assert_eq!( + rows(&batches), + expected + .iter() + .map(|(value, label)| (*value, label.map(str::to_string))) + .collect::>() + ); + } +} + +#[tokio::test] +async fn test_direct_or_match_modifiers() { + for (modifier, left, right, expected) in [ + (None, "left", "right", 2), + (or_modifier("lhs or on(k) rhs"), "same", "same", 1), + (or_modifier("lhs or on() rhs"), "left", "right", 1), + (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), + ] { + let (_, batches) = run( + &matrix_source("lhs", Some(Some(left)), 1, 1.0), + &matrix_source("rhs", Some(Some(right)), 1, 2.0), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &modifier, + ) + .await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected + ); + } +} + +#[tokio::test] +async fn test_direct_or_nested_projection_uses_left_context() { + let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); + let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert!(raw.schema().iter().any(|(qualifier, field)| { + qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" + })); + let nested = LogicalPlanBuilder::from(raw) + .project(vec![ + DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(DfExpr::Column(Column::new( + Some(TableReference::bare("lhs")), + "v", + ))), + op: Operator::Plus, + right: Box::new(lit(1.0)), + }) + .alias("v_plus"), + ]) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(nested, &build_query_engine_state()).await; + assert_eq!(values(&batches, "v_plus"), vec![2.0]); +} + +#[tokio::test] +async fn test_direct_or_skips_user_internal_key_name() { + const USER_TAG: &str = "__promql_or_match_0"; + let left = tagged_source( + "lhs", + false, + (USER_TAG, Some("left")), + DirectOrValue::Float64(1.0), + ); + let right = tagged_source( + "rhs", + false, + (USER_TAG, Some("right")), + DirectOrValue::Float64(2.0), + ); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", USER_TAG], "v"), + direct_or_context("rhs", &["job", USER_TAG], "v"), + &or_modifier("lhs or on(missing_label) rhs"), + ) + .await; + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_1@") + ); + let (_, batches) = execute(raw, &build_query_engine_state()).await; + assert!( + batches + .iter() + .all(|batch| batch.column_by_name(USER_TAG).is_some()) + ); +} + +#[tokio::test] +async fn test_direct_or_substrait_round_trip_with_normalized_key() { + let state = build_query_engine_state(); + let ctx = SessionContext::new_with_state(state.session_state()); + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog + .register_schema("public", Arc::new(MemorySchemaProvider::new())) + .unwrap(); + ctx.register_catalog("datafusion", catalog); + let left = matrix_source("lhs", Some(Some("")), 1, 1.0); + let right = matrix_source("rhs", None, 1, 2.0); + ctx.register_table( + TableReference::full("datafusion", "public", "lhs"), + table(&left), + ) + .unwrap(); + ctx.register_table( + TableReference::full("datafusion", "public", "rhs"), + table(&right), + ) + .unwrap(); + let raw = plan_direct_or( + ctx.table("datafusion.public.lhs") + .await + .unwrap() + .into_unoptimized_plan(), + ctx.table("datafusion.public.rhs") + .await + .unwrap() + .into_unoptimized_plan(), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + let decoded = DFLogicalSubstraitConvertor + .decode( + DFLogicalSubstraitConvertor + .encode(&raw, DefaultSerializer) + .unwrap(), + ctx.state(), + ) + .await + .unwrap(); + let (optimized, batches) = execute(decoded, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); + assert_eq!(values(&batches, "v"), vec![1.0]); +} + +#[tokio::test] +async fn test_direct_or_numeric_value_types() { + let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); + let right = tagged_source( + "rhs", + false, + ("k", Some("rhs")), + DirectOrValue::Float64(0.5), + ); + let (optimized, batches) = run( + &left, + &right, + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert_eq!( + optimized + .schema() + .field_with_name(None, "v") + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!(values(&batches, "v"), vec![0.5]); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider: provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job"], "v"); + let right_context = direct_or_context("rhs", &["job"], "v"); + let error = planner + .or_operator( + scan(&job_source("lhs", DirectOrValue::Utf8("x"))), + scan(&job_source("rhs", DirectOrValue::Float64(1.0))), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on() rhs"), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("OR value fields have incompatible types") + ); +} + +#[tokio::test] +async fn test_or_with_histogram_quantile_missing_le_column() { + let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; + let eval_stmt = build_eval_stmt(case); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + assert_normal_metric_schema(&plan); +} + +#[tokio::test] +async fn test_or_with_right_empty_histogram_restores_left_context() { + let eval_stmt = build_eval_stmt( + r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#, + ); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_or_with_both_empty_histograms() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#, + ); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "left_histogram_bucket".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "right_histogram_bucket".to_string(), + ), + ], + &["pod", "instance"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + match plan { + LogicalPlan::EmptyRelation(relation) => { + assert!(!relation.produce_one_row); + assert!(!relation.schema.fields().is_empty()); + assert!( + relation + .schema + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64) + ); + assert!( + relation + .schema + .fields() + .iter() + .any(|field| field.name() == "pod") + ); + assert!( + !relation + .schema + .fields() + .iter() + .any(|field| field.name() == "instance") + ); + } + _ => panic!("Expected EmptyRelation, but got: {plan:?}"), + } +} + +#[tokio::test] +async fn test_nested_or_with_both_empty_histograms() { + for case in [ + r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#, + r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#, + ] { + let eval_stmt = build_eval_stmt(case); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "left_histogram_bucket".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "right_histogram_bucket".to_string(), + ), + ], + &["pod", "instance"], + ) + .await; + + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn test_or_with_empty_histogram_modifiers() { + for case in [ + r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#, + r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#, + ] { + let eval_stmt = build_eval_stmt(case); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + assert_normal_metric_schema(&plan); + } +} + +#[tokio::test] +async fn test_unless_preserves_left_context_for_histogram() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &state, + ) + .await + .unwrap(); + assert!(contains_histogram_fold(&plan), "{plan:?}"); + let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await; + assert!(contains_histogram_fold(&optimized), "{optimized:?}"); + let batches = datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + assert!(batches.iter().all(|batch| batch.num_rows() == 0)); +} + +#[tokio::test] +async fn test_and_preserves_left_context_for_histogram() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#, + ); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!(contains_histogram_fold(&plan), "{plan:?}"); +} + +#[tokio::test] +async fn test_and_preserves_left_context_when_le_is_missing() { + let eval_stmt = + build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}"); + assert!(!plan.schema().fields().is_empty()); + assert!(!contains_histogram_fold(&plan), "{plan:?}"); +} + +async fn build_matching_filter_plan(query: &str) -> String { + let table_provider = build_test_table_provider_with_distinct_tags(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() +} + +/// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value +/// field that is neither a primary key nor a value column of the metric, so +/// `count by(status) (...)` still reports it among the aggregation's tag columns. +async fn build_test_table_provider_with_string_field( + table_tags: &[(&str, &[&str])], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (table_name, tags) in table_tags { + let mut columns = tags + .iter() + .map(|tag| { + ColumnSchema::new( + (*tag).to_string(), + ConcreteDataType::string_datatype(), + false, + ) + }) + .collect::>(); + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + columns.push(ColumnSchema::new( + "status".to_string(), + ConcreteDataType::string_datatype(), + true, + )); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(columns))) + .primary_key_indices((0..tags.len()).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name((*table_name).to_string()) + .meta(table_meta) + .build() + .unwrap(); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: (*table_name).to_string(), + table_id: 1024, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_matching_filter_plan_with_string_field(query: &str) -> String { + let table_provider = build_test_table_provider_with_string_field(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_both_operands() { + for query in [ + r#"metric_a / metric_b{host="foo"}"#, + r#"metric_a / on(host, device) metric_b{host="foo"}"#, + r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#, + r#"metric_a / ignoring(device) metric_b{host="foo"}"#, + r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#, + r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() { + for query in [ + r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#, + r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#, + r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } + // A global ranking one-side must see every host, so the matcher stays put. + let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 1, + "{query}\n{plan}" + ); +} + +#[tokio::test] +async fn binary_matching_label_filter_skips_selecting_aggregations() { + // `topk` ranks its input, so filtering before it changes the candidate set. + let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); +} + +#[tokio::test] +async fn binary_value_field_matcher_stays_on_its_own_operand() { + let value = greptime_value(); + for query in [ + format!(r#"metric_a / metric_b{{{value}="2"}}"#), + format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#), + ] { + let plan = build_matching_filter_plan(&query).await; + assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}"); + } +} + +#[tokio::test] +async fn binary_value_field_matcher_is_not_copied_across_aggregations() { + // `status` varies between the samples of one series, so filtering the other operand by it + // would drop the newest sample before sample selection (#9242). + for query in [ + r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#, + r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + ] { + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!( + plan.matches(r#"Utf8("ready")"#).count(), + 1, + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() { + // `count`, not `sum`: the string value field is not summable. + let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#; + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}"); +} + +#[tokio::test] +async fn binary_matching_label_filter_skips_unproven_expressions() { + for query in [ + r#"metric_a > on(host, device) metric_b{host="foo"}"#, + r#"metric_a / on(host) group_left metric_b{host="foo"}"#, + r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#, + r#"metric_a / on(device) metric_b{host="foo"}"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); + } +} + +#[tokio::test] +async fn test_or_context_uses_left_qualified_output() { + let case = r#"(normal_metric or other_metric) + 1"#; + let eval_stmt = build_eval_stmt(case); + let state = build_query_engine_state(); + let plan = + PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64), + "{plan:?}" + ); + let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await; +} + +#[tokio::test] +async fn test_or_context_uses_left_qualified_empty_histogram_output() { + let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#; + let eval_stmt = build_eval_stmt(case); + let plan = PromPlanner::stmt_to_plan( + build_or_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64), + "{plan:?}" + ); +} + +#[tokio::test] +async fn test_direct_or_preserves_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (planner, plan) = mixed_direct_or(histogram_on_left).await; + + let float_field = &planner.ctx.field_columns[0]; + let histogram_field = &planner.ctx.field_columns[1]; + assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); + assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); + assert_eq!( + plan.schema() + .field_with_name(None, float_field) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + plan.schema() + .field_with_name(None, histogram_field) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + + let (optimized, batches) = execute(plan, &build_query_engine_state()).await; + assert_no_internal_or_keys(optimized.schema()); + let mut sample_kinds = batches + .iter() + .flat_map(|batch| { + let values = batch.column_by_name(float_field).unwrap(); + let histograms = batch.column_by_name(histogram_field).unwrap(); + (0..batch.num_rows()).map(|row| (values.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + sample_kinds.sort_unstable(); + assert_eq!(sample_kinds, vec![(false, true), (true, false)]); + } +} + +#[tokio::test] +async fn malformed_classic_bucket_does_not_drop_native_histogram() { + let state = build_query_engine_state(); + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &value_field), vec![0.0]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert!(infos.is_empty()); +} + +#[tokio::test] +async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!( + plan_text.contains("prom_native_histogram_mul_scalar"), + "{plan_text}" + ); + assert!( + plan_text.contains("prom_native_histogram_scalar_mul"), + "{plan_text}" + ); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert!(values(&batches, &float_field).is_empty()); + let mut sums = histograms(&batches, &histogram_field) + .into_iter() + .map(|histogram| histogram.sum) + .collect::>(); + sums.sort_by(f64::total_cmp); + assert_eq!(sums, vec![2.0, 3.0]); +} + +#[tokio::test] +async fn test_mixed_binary_operator_reports_only_dropped_samples() { + for (query, expected_rows, expected_infos) in [ + ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), + ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), + ("(lf or on(tag) lh) % on(tag) lh", 0, 1), + ] { + let state = build_query_engine_state(); + let annotations = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + Some(annotations.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows, + "{query}" + ); + let mut warnings = vec![]; + let mut infos = vec![]; + annotations.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty(), "{query}: {warnings:?}"); + assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); + } +} + +#[tokio::test] +async fn test_histogram_only_min_drops_empty_aggregate_group() { + // `min` over native-histogram-only input drops every sample in the group, so the + // NULL-valued aggregate row must be filtered out. Otherwise an outer expression + // like `group()` resurrects the group Prometheus considers unseen. + let state = build_query_engine_state(); + for query in ["min(lh)", "group(min(lh))"] { + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "{query}" + ); + } +} + +#[tokio::test] +async fn test_mixed_min_drops_histogram_only_group() { + // With alternative float/histogram fields, `min by (tag)` keeps float-only groups + // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of + // emitting a NULL-valued row for them. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &float_field), vec![2.0]); +} + +#[tokio::test] +async fn test_mixed_or_can_feed_another_or() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + let mut float_values = values(&batches, &float_field); + float_values.sort_by(f64::total_cmp); + assert_eq!(float_values, vec![2.0, 7.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); +} + +#[tokio::test] +async fn test_mixed_fields_align_with_single_float_vector() { + let (planner, mixed) = mixed_direct_or(false).await; + let scale = tagged_source( + "scale", + false, + ("k", Some("float")), + DirectOrValue::Float64(2.0), + ); + let scale = scan(&scale); + let scale_fields = vec!["v".to_string()]; + let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { + unreachable!() + }; + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + mixed.schema(), + scale.schema(), + &planner.ctx.field_columns, + &scale_fields, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(_, right)| *right == &scale_fields[0]) + ); + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + scale.schema(), + mixed.schema(), + &scale_fields, + &planner.ctx.field_columns, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(left, _)| *left == &scale_fields[0]) + ); +} + +#[tokio::test] +async fn test_non_bool_comparison_filters_mixed_sample_lanes() { + let (planner, input) = mixed_direct_or(false).await; + let input_schema = input.schema().clone(); + let plan = planner + .filter_on_field_column(input, |field| { + if PromPlanner::field_column_is_native_histogram(&input_schema, field) { + Ok(lit(false)) + } else { + Ok(col(field).gt(lit(0.0))) + } + }) + .unwrap(); + let float_field = planner.ctx.field_columns[0].clone(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_field), vec![1.25]); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); +} + +#[tokio::test] +async fn test_mixed_left_and_unless_preserve_sample_lanes() { + for (expression, expected_sample_kind) in [ + ("lhs and on(k) mask", (false, true)), + ("lhs unless on(k) mask", (true, false)), + ] { + let (mut planner, left) = mixed_direct_or(false).await; + let left_context = planner.ctx.clone(); + let float_field = left_context.field_columns[0].clone(); + let histogram_field = left_context.field_columns[1].clone(); + let mask = tagged_source( + "mask", + false, + ("k", Some("histogram")), + DirectOrValue::Float64(1.0), + ); + let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { + unreachable!() + }; + let plan = planner + .set_op_on_non_field_columns( + left, + scan(&mask), + left_context, + direct_or_context("mask", &["job", "k"], "v"), + binary.op, + &binary.modifier, + ) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let sample_kinds = batches + .iter() + .flat_map(|batch| { + let floats = batch.column_by_name(&float_field).unwrap(); + let histograms = batch.column_by_name(&histogram_field).unwrap(); + (0..batch.num_rows()).map(|row| (floats.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); + } +} + +#[tokio::test] +async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("some_metric * scalar(vector(2))"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let schema = plan.schema(); + assert_eq!( + schema + .field_with_unqualified_name(greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + schema + .field_with_unqualified_name(greptime_native_histogram()) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + assert!( + plan.display_indent_schema() + .to_string() + .contains("prom_native_histogram_mul_scalar"), + "{plan:?}" + ); +} + +#[tokio::test] +async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![0.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); +} + +#[tokio::test] +async fn test_unary_negates_mixed_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (mut planner, input) = mixed_direct_or(histogram_on_left).await; + let plan = planner.negate_field_columns(input).unwrap(); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let float_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap(); + let histogram_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, float_field), vec![-1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(histogram_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, -1.0); + assert_eq!(histogram.sum, -1.0); + assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); + } +} + +#[tokio::test] +async fn test_native_histogram_sum_and_avg_execute_real_batches() { + for op_name in ["sum", "avg"] { + for incompatible in [false, true] { + let mut second = direct_or_histogram(); + if incompatible { + second.schema = CUSTOM_BUCKETS_SCHEMA; + second.custom_values = vec![1.0]; + } + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = + mixed_aggregate_input(vec![direct_or_histogram(), second]).await; + planner.promql_annotations = Some(collector.clone()); + let histogram_column = planner.ctx.field_columns[1].clone(); + planner.ctx.field_columns = vec![histogram_column.clone()]; + let input = LogicalPlanBuilder::from(input) + .project([col("ts"), col(&histogram_column)]) + .unwrap() + .build() + .unwrap(); + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse(&format!("{op_name}(mixed)")).unwrap() + else { + unreachable!() + }; + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter(planner.create_empty_values_filter_expr(false).unwrap()) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(infos.is_empty()); + if incompatible { + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + assert!(warnings.iter().any(|warning| { + warning + == &format!( + "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" + ) + })); + } else { + let histograms = histograms(&batches, &histogram_column); + assert_eq!(histograms.len(), 1); + let expected = if op_name == "sum" { 2.0 } else { 1.0 }; + assert_eq!(histograms[0].count, expected); + assert_eq!(histograms[0].sum, expected); + assert!(warnings.is_empty()); + } + } + } +} + +#[tokio::test] +async fn test_canonical_mixed_count_group_and_count_values_execute() { + let state = build_query_engine_state(); + for (query, expected) in [ + ("count(some_metric)", vec![2.0]), + ("group(some_metric)", vec![1.0]), + (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_sample_count")), + "{query}: {plan:?}" + ); + let value_fields = plan + .schema() + .fields() + .iter() + .filter(|field| { + matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) || field.data_type() == &native_histogram_value_type().as_arrow_type() + }) + .collect::>(); + assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); + assert_ne!( + value_fields[0].data_type(), + &native_histogram_value_type().as_arrow_type(), + "{query}: {plan:?}" + ); + let value_column = value_fields[0].name().clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = numeric_values(&batches, &value_column); + actual.sort_by(f64::total_cmp); + assert_eq!(actual, expected, "{query}"); + + if query.starts_with("count_values") { + let mut sample_labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("sample") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(str::to_string) + }) + .collect::>(); + sample_labels.sort(); + let mut expected_labels = vec!["2".to_string(), direct_or_histogram().promql_string()]; + expected_labels.sort(); + assert_eq!(sample_labels, expected_labels); + } + } +} + +#[tokio::test] +async fn test_mixed_or_sum_aggregates_each_sample_type() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts"), col("k")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([ + col(&float_column), + col(&histogram_column), + col("ts"), + col("k"), + ]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(&histogram_column)? + .as_any() + .downcast_ref::()?; + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, 1.0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!( + warnings, + vec!["sum: dropped aggregation result containing both float and native histogram samples"] + ); +} + +#[tokio::test] +async fn test_mixed_or_sum_drops_incompatible_mixed_group() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + let mut custom = direct_or_histogram(); + custom.schema = CUSTOM_BUCKETS_SCHEMA; + custom.custom_values = vec![1.0]; + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.iter().any(|warning| { + warning + == "sum: dropped aggregation result containing both float and native histogram samples" + })); +} + +#[tokio::test] +async fn test_mixed_or_min_records_only_present_histograms() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("min(lhs)").unwrap() + else { + unreachable!() + }; + let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; + + for (histograms, expected_infos) in [ + (vec![], vec![]), + (vec![direct_or_histogram()], vec![expected_info]), + ] { + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(histograms).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_ignored_histogram_filter_expr(op, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!(infos, expected_infos); + } +} + +#[tokio::test] +async fn test_mixed_or_value_aliases_do_not_replace_labels() { + let left = source( + "lhs", + false, + 1, + vec![("job", Some("job")), ("k", Some("float"))], + DirectOrValue::Float64(1.0), + ); + let right = source( + "rhs", + false, + 1, + vec![ + ("job", Some("job")), + ("k", Some("histogram")), + (greptime_value(), Some("value-label")), + ], + DirectOrValue::NativeHistogram(direct_or_histogram()), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left = LogicalPlanBuilder::from(scan(&left)) + .project(vec![ + col("ts"), + col("job"), + col("k"), + col("v").alias(greptime_value()), + ]) + .unwrap() + .build() + .unwrap(); + let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); + let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); + let plan = planner + .or_operator( + left, + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + + assert_eq!( + plan.schema() + .field_with_name(None, greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Utf8 + ); + assert!( + planner + .ctx + .field_columns + .iter() + .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) + ); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(greptime_value()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect::>(); + assert_eq!(labels, vec!["value-label"]); +} + +#[tokio::test] +async fn test_mixed_or_routes_float_histogram_and_label_functions() { + for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { + unreachable!() + }; + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state, None) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let values = batches + .iter() + .flat_map(|batch| { + batch + .schema() + .fields() + .iter() + .position(|field| field.data_type() == &ArrowDataType::Float64) + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .into_iter() + .flatten() + }) + .collect::>(); + assert_eq!(values, vec![expected], "{function}"); + } + + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = + parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() + else { + unreachable!() + }; + let args = planner.create_function_args(&call.args.args).unwrap(); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, args.literals, input.schema(), &state, None) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); + assert_eq!(sample_count, 2); +} + +/// Table provider with a single metric `cv_metric` holding three series at `ts=1000`: +/// `k="a"` and `k="c"` carry the value 1.0, `k="b"` carries 2.0. +async fn build_count_values_table_provider() -> DfTableSourceProvider { + build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await +} + +/// Table provider with a single metric `cv_metric` holding one series per given value at +/// `ts=1000` (`k` names the series, the sample value is the given value). +async fn build_count_values_table_provider_with_values(values: &[f64]) -> DfTableSourceProvider { + build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from( + values.to_vec(), + ))) + .await +} + +/// Like [`build_count_values_table_provider_with_values`], but with a caller provided value +/// column, so tests can cover value columns that are not `Float64` (e.g. `BIGINT`). +async fn build_count_values_table_provider_with_value_array( + values: ArrayRef, +) -> DfTableSourceProvider { + let value_data_type = ConcreteDataType::from_arrow_type(values.data_type()); + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new(greptime_value().to_string(), value_data_type, true), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(3_001) + .name("cv_metric") + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from( + (0..values.len()) + .map(|index| format!("k{index}")) + .collect::>(), + )), + Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])), + values.clone(), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + "cv_metric", + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 3_001, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "cv_metric".to_string(), + table_id: 3_001, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +/// Collects `(label, value)` pairs of a `count_values` result, where `label` is the +/// PromQL label generated by `count_values` and `value` is the aggregated sample value. +/// +/// The generated label holds the original sample value in PromQL's textual form, so it +/// is asserted as a string: comparing it as a number would not catch formatting bugs +/// (`1.0` instead of `1`, scientific notation, ...). +fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + // The aggregated value is the only numeric column that is not the generated label. + let value_index = batch + .schema() + .fields() + .iter() + .position(|field| { + field.name() != label + && matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) + }) + .expect("no aggregated value column"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + let values = datafusion::arrow::compute::cast( + batch.column(value_index), + &ArrowDataType::Float64, + ) + .unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + labels + .iter() + .zip(values.iter()) + .map(|(label, value)| (label.unwrap(), value.unwrap())) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1))); + rows +} + +/// Asserts that a `count_values` result holds one sample per label set and evaluation +/// timestamp: Prometheus groups by the generated label, so a timestamp must never repeat +/// a label set (that would mean the samples were still grouped by the overwritten input +/// label). +fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) { + let mut seen = HashMap::>::new(); + for batch in batches { + let timestamp_index = batch + .schema() + .fields() + .iter() + .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) + .expect("no timestamp column"); + let timestamps = batch + .column(timestamp_index) + .as_any() + .downcast_ref::() + .expect("timestamp column is not a millisecond timestamp"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + for (timestamp, label) in timestamps.iter().zip(labels.iter()) { + let timestamp = timestamp.unwrap(); + let label = label.unwrap(); + assert!( + seen.entry(timestamp).or_default().insert(label.to_string()), + "duplicated label set `{label}` at timestamp {timestamp}" + ); + } + } +} + +#[tokio::test] +async fn test_count_values_generated_label_survives_enclosing_expr() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + for (case, label) in [ + (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"), + ( + r#"abs(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"round(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"), + ( + r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#, + "v", + ), + ( + r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#, + "v", + ), + // The generated label overwrites an input label with the same name. + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#, + "ip", + ), + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#, + "ip", + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + )], + &["ip"], + ) + .await, + &build_eval_stmt(case), + &build_query_engine_state(), + ) + .await + .unwrap(); + + let label_columns = plan + .schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(); + assert_eq!( + label_columns, + 1, + "{case}: the `{label}` label must survive: {}", + plan.display_indent() + ); + } +} + +#[tokio::test] +async fn test_count_values_generated_label_in_enclosing_expr_execute() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + let state = build_query_engine_state(); + // (query, generated label, expected `(label value, aggregated value)` pairs) + for (query, label, expected) in [ + ( + r#"count_values("v", cv_metric)"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"abs(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"round(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"count_values("v", cv_metric) + 1"#, + "v", + vec![("1", 3.0), ("2", 2.0)], + ), + ( + r#"sum by (v) (count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"topk(10, count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + // The generated label overwrites the input label with the same name, and the + // samples are grouped by the generated label only: `{k="1"}` holds the two + // samples of value `1.0` instead of one row per (overwritten label, value). + ( + r#"count_values("k", cv_metric) by (k)"#, + "k", + vec![("1", 2.0), ("2", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider().await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap_or_else(|err| panic!("{query}: {err}")); + + assert_eq!( + plan.schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(), + 1, + "{query}: {}", + plan.display_indent() + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(count_values_rows(&batches, label), expected, "{query}"); + assert_unique_label_set_per_timestamp(&batches, label); + } +} + +#[tokio::test] +async fn test_count_values_label_is_prometheus_formatted_value() { + // PromQL materializes the generated label with `strconv.FormatFloat(value, 'f', -1, 64)`: + // the shortest decimal form of the sample value without an exponent. The label is a + // label, so it must be a string column holding exactly that text: arrow's + // `Float64 -> Utf8` cast would render `1`/`200`/`1e21` as `1.0`/`200.0`/`1e21`. + let state = build_query_engine_state(); + // Samples are grouped by that formatted text, exactly like Prometheus groups by the + // generated label: `-0.0` and `0.0` are two series (`-0` and `0`), while values that + // round to the same text share one group. `0.0` formats as "0". + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_values(&[ + -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5, + ]) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_unique_label_set_per_timestamp(&batches, "v"); + let mut labels = count_values_rows(&batches, "v") + .into_iter() + .map(|(label, _)| label) + .collect::>(); + labels.sort(); + assert_eq!( + labels, + vec![ + "-0", + "0", + "0.0000001", + "0.5", + "1", + "1000000000000000000000", + "2.5", + "200", + ] + ); +} + +#[tokio::test] +async fn test_count_values_groups_by_formatted_value_for_bigint_input() { + // The grouping key of `count_values` is the formatted sample value, not the raw input + // value. Two `BIGINT` values that differ below the `Float64` precision (`2^53` and + // `2^53 + 1`) cast and format to the same label, so they must share one group and one + // count, exactly like Prometheus, which groups by the generated label text. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![ + 9_007_199_254_740_992_i64, + 9_007_199_254_740_993_i64, + ]))) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + // One timestamp must never carry the same label set twice. + assert_unique_label_set_per_timestamp(&batches, "v"); + assert_eq!( + count_values_rows(&batches, "v"), + vec![("9007199254740992", 2.0)] + ); +} + +#[tokio::test] +async fn eval_times_before_epoch() { + let plan = |start: SystemTime, end: SystemTime| async move { + let mut eval_stmt = build_eval_stmt("some_metric"); + eval_stmt.start = start; + eval_stmt.end = end; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()).await + }; + + plan( + UNIX_EPOCH - Duration::from_secs(30_000), + UNIX_EPOCH - Duration::from_secs(29_000), + ) + .await + .unwrap(); + // Beyond the i64 millisecond range, which must not wrap into the future. Windows cannot + // represent this SystemTime. + if let Some(too_early) = UNIX_EPOCH.checked_sub(Duration::from_secs(10_000_000_000_000_000)) { + assert!(plan(too_early, too_early).await.is_err()); + } +} diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index 3727ea9a30c..9ca8a8b646c 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -475,7 +475,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() { ); let state = build_query_engine_state(); let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) + .create_function_expr(&call.func, vec![], input.schema(), &state, None) .unwrap(); exprs.insert(0, planner.create_time_index_column_expr().unwrap()); let plan = LogicalPlanBuilder::from(input) @@ -535,6 +535,7 @@ async fn matchers_read_absent_labels_as_empty() { schema.clone(), vec![Arc::new(StringArray::from(vec![ None, + Some(""), Some(GREPTIME_TEMPORALITY_DELTA), ]))], ) @@ -542,10 +543,19 @@ async fn matchers_read_absent_labels_as_empty() { let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); for (matcher, expected) in [ - (r#"="""#, vec![None]), - (r#"!="delta""#, vec![None]), - (r#"=~".*""#, vec![None, Some(GREPTIME_TEMPORALITY_DELTA)]), - (r#"!~"delta""#, vec![None]), + (r#"="""#, vec![None, Some("")]), + (r#"!="delta""#, vec![None, Some("")]), + ( + r#"=~".*""#, + vec![None, Some(""), Some(GREPTIME_TEMPORALITY_DELTA)], + ), + (r#"!~"delta""#, vec![None, Some("")]), + (r#"!="""#, vec![Some(GREPTIME_TEMPORALITY_DELTA)]), + ( + r#"=~"delta|""#, + vec![None, Some(""), Some(GREPTIME_TEMPORALITY_DELTA)], + ), + (r#"!~"delta|""#, vec![]), (r#"="delta""#, vec![Some(GREPTIME_TEMPORALITY_DELTA)]), ] { let query = format!(r#"metric{{{marker}{matcher}}}"#); @@ -558,10 +568,10 @@ async fn matchers_read_absent_labels_as_empty() { }; let expressions = PromPlanner::matchers_to_expr(selector.matchers, scan.schema()).unwrap(); let display = expressions.iter().map(ToString::to_string).join(" AND "); - if matcher == r#"="delta""# { - assert!(!display.contains("coalesce"), "{display}"); + if [r#"="delta""#, r#"!="""#, r#"!~"delta|""#].contains(&matcher) { + assert!(!display.contains("IS NULL"), "{display}"); } else if !expressions.is_empty() { - assert!(display.contains("coalesce"), "{display}"); + assert!(display.contains("IS NULL"), "{display}"); } let plan = if let Some(filter) = conjunction(expressions) { LogicalPlanBuilder::from(scan) @@ -599,7 +609,7 @@ async fn matchers_read_absent_labels_as_empty() { // The rule is about NULL, not about the marker: any nullable label column // can be NULL where the series does not carry the label. A non-nullable one // has nothing to normalize. - for (nullable, wants_coalesce) in [(true, true), (false, false)] { + for (nullable, wants_null_filter) in [(true, true), (false, false)] { let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new( "label", ArrowDataType::Utf8, @@ -626,8 +636,8 @@ async fn matchers_read_absent_labels_as_empty() { .map(ToString::to_string) .join(" AND "); assert_eq!( - wants_coalesce, - expressions.contains("coalesce"), + wants_null_filter, + expressions.contains("IS NULL"), "nullable={nullable}: {expressions}" ); } @@ -1102,12 +1112,12 @@ async fn delta_offsets_survive_optimized_plan_serialization() { ( "timestamp positive offset", r#"timestamp(delta_metric{series="cumulative"} offset 60s)"#, - 120.0, + 60.0, ), ( "timestamp negative offset", r#"timestamp(delta_metric{series="cumulative"} offset -60s)"#, - 120.0, + 180.0, ), ( "range positive offset", diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index d6904c7ece0..845bc44437f 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -14,19 +14,21 @@ use std::sync::Arc; +use bytes::Bytes; +use catalog::CatalogManagerRef; use common_error::ext::BoxedError; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; use common_query::error::RegisterUdfSnafu; use common_query::logical_plan::SubstraitPlanDecoder; -use datafusion::catalog::CatalogProviderList; -use datafusion::common::DataFusionError; +use datafusion::catalog::{CatalogProviderList, TableProvider}; +use datafusion::common::{DFSchema, DataFusionError, TableReference, not_impl_err, substrait_err}; use datafusion::error::Result; use datafusion::execution::context::SessionState; use datafusion::execution::registry::SerializerRegistry; use datafusion::execution::{FunctionRegistry, SessionStateBuilder}; -use datafusion::logical_expr::LogicalPlan; -use datafusion_expr::UserDefinedLogicalNode; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion_expr::{Expr, UserDefinedLogicalNode}; use greptime_proto::substrait_extension::MergeScan as PbMergeScan; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, @@ -44,12 +46,18 @@ use promql::functions::{ Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use prost::Message; -use session::context::QueryContextRef; +use session::context::{QueryContext, QueryContextRef}; use snafu::ResultExt; +use substrait::df_logical_plan::consumer::{DefaultSubstraitConsumer, SubstraitConsumer}; +use substrait::error::{DecodeDfPlanSnafu, DecodeRelSnafu}; use substrait::extension_serializer::ExtensionSerializer; -use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; +use substrait::substrait_proto_df::proto::{ + ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, Plan, Type, +}; +use substrait::{DFLogicalSubstraitConvertor, Extensions, SubstraitPlan}; use crate::dist_plan::MergeScanLogicalPlan; +use crate::query_engine::QueryEngineState; /// Extended [`substrait::extension_serializer::ExtensionSerializer`] but supports [`MergeScanLogicalPlan`] serialization. #[derive(Debug)] @@ -86,8 +94,7 @@ impl SerializerRegistry for DefaultSerializer { bytes: &[u8], ) -> Result> { if name == MergeScanLogicalPlan::name() { - // TODO(dennis): missing `session_state` to decode the logical plan in `MergeScanLogicalPlan`, - // so we only save the unoptimized logical plan for view currently. + // This registry has no session state; the query engine decodes `MergeScan` payloads. Err(DataFusionError::Substrait(format!( "Unsupported plan node: {name}" ))) @@ -97,6 +104,196 @@ impl SerializerRegistry for DefaultSerializer { } } +/// Async `MergeScan` payload decoding over DataFusion's default Substrait consumer. +struct MergeScanSubstraitConsumer<'a> { + inner: DefaultSubstraitConsumer<'a>, + session_state: &'a SessionState, + catalog_manager: Option, + payload: bool, +} + +impl MergeScanSubstraitConsumer<'_> { + fn query_ctx(&self) -> QueryContextRef { + self.session_state + .config() + .get_extension() + .unwrap_or_else(QueryContext::arc) + } + + async fn decode_payload(&self, payload: Vec) -> Result { + decode_plan(Bytes::from(payload), self.session_state, true) + .await + .map_err(|e| DataFusionError::External(Box::new(e))) + } +} + +#[async_trait::async_trait] +impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { + async fn resolve_table_ref( + &self, + table_ref: &TableReference, + ) -> Result>> { + if !self.payload { + return self.inner.resolve_table_ref(table_ref).await; + } + + let query_ctx = self.query_ctx(); + let resolved = table_ref + .clone() + .resolve(query_ctx.current_catalog(), &query_ctx.current_schema()); + // Manual states retain DataFusion's own catalog and information-schema behavior. + let Some(manager) = &self.catalog_manager else { + return self + .inner + .resolve_table_ref(&TableReference::from(resolved)) + .await; + }; + let catalog_list: Arc = Arc::new( + catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( + manager.clone(), + query_ctx, + ), + ); + let schema_name = resolved.schema.as_ref(); + let catalog_name = resolved.catalog.as_ref(); + let schema = catalog_list + .catalog(catalog_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve catalog: {catalog_name}")) + })? + .schema(schema_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve schema: {schema_name}")) + })?; + schema.table(&resolved.table).await + } + + fn get_extensions(&self) -> &Extensions { + self.inner.get_extensions() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + self.inner.get_function_registry() + } + + fn push_outer_schema(&self, schema: Arc) { + self.inner.push_outer_schema(schema) + } + + fn pop_outer_schema(&self) { + self.inner.pop_outer_schema() + } + + fn get_outer_schema(&self, steps_out: usize) -> Option> { + self.inner.get_outer_schema(steps_out) + } + + fn push_lambda_parameters( + &self, + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> Result> { + self.inner + .push_lambda_parameters(lambda_parameters, input_schema) + } + + fn pop_lambda_parameters(&self) { + self.inner.pop_lambda_parameters() + } + + fn lambda_variable(&self, steps_out: usize, field_idx: usize) -> Result { + self.inner.lambda_variable(steps_out, field_idx) + } + + async fn consume_extension_leaf(&self, rel: &ExtensionLeafRel) -> Result { + // Unknown and detail-less leaves keep the default consumer's errors. + let Some(detail) = rel.detail.as_ref() else { + return self.inner.consume_extension_leaf(rel).await; + }; + if detail.type_url != MergeScanLogicalPlan::name() { + return self.inner.consume_extension_leaf(rel).await; + } + + let merge_scan = PbMergeScan::decode(detail.value.as_ref()).map_err(|e| { + DataFusionError::Substrait(format!("Failed to decode the MergeScan plan node: {e}")) + })?; + + let input = self.decode_payload(merge_scan.input).await?; + + // `PbMergeScan` lacks `partition_cols`; decoded plans lose this optimization and may repartition. + Ok( + MergeScanLogicalPlan::new(input, merge_scan.is_placeholder, Default::default()) + .into_logical_plan(), + ) + } + + // Keep child dispatch on this consumer; delegating the parent bypasses async MergeScan decoding. + async fn consume_extension_single(&self, rel: &ExtensionSingleRel) -> Result { + let Some(detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionSingleRel"); + }; + let plan = self + .session_state + .serializer_registry() + .deserialize_logical_plan(&detail.type_url, &detail.value)?; + let Some(input_rel) = &rel.input else { + return substrait_err!( + "ExtensionSingleRel missing input rel, try using ExtensionLeafRel instead" + ); + }; + let input_plan = self.consume_rel(input_rel).await?; + let plan = plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + } + + // Multi-input extension children likewise need this consumer's dispatch. + async fn consume_extension_multi(&self, rel: &ExtensionMultiRel) -> Result { + let Some(detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionMultiRel"); + }; + let plan = self + .session_state + .serializer_registry() + .deserialize_logical_plan(&detail.type_url, &detail.value)?; + let mut inputs = Vec::with_capacity(rel.inputs.len()); + for input in &rel.inputs { + let input_plan = self.consume_rel(input).await?; + inputs.push(input_plan); + } + let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + } +} + +/// Decode with async MergeScan handling, preserving the default consumer's extension checks. +async fn decode_plan( + bytes: Bytes, + state: &SessionState, + payload: bool, +) -> substrait::error::Result { + let plan = Plan::decode(bytes).context(DecodeRelSnafu)?; + let extensions = Extensions::try_from(&plan.extensions).context(DecodeDfPlanSnafu)?; + if !extensions.type_variations.is_empty() { + return not_impl_err!("Type variation extensions are not supported") + .context(DecodeDfPlanSnafu); + } + + let catalog_manager = state + .config() + .get_extension::() + .map(|engine_state| engine_state.catalog_manager().clone()); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, state), + session_state: state, + catalog_manager, + payload, + }; + + DFLogicalSubstraitConvertor + .decode_with_consumer(&plan, state, &consumer) + .await +} + /// The datafusion `[LogicalPlan]` decoder. pub struct DefaultPlanDecoder { session_state: SessionState, @@ -115,6 +312,108 @@ impl DefaultPlanDecoder { } } +/// Registers Greptime functions after building the request decoder state. +/// +/// DataFusion aliases can rebind Greptime `date_format` to built-in `to_char`; registration must +/// happen after `SessionStateBuilder::build` on the state that decodes the request plan. +fn register_greptime_functions( + session_state: &mut SessionState, + query_ctx: &QueryContextRef, +) -> common_query::error::Result<()> { + for func in FUNCTION_REGISTRY.scalar_functions() { + let udf = func.provide(FunctionContext { + query_ctx: query_ctx.clone(), + state: Default::default(), + }); + session_state + .register_udf(Arc::new(udf)) + .context(RegisterUdfSnafu { name: func.name() })?; + } + + for func in FUNCTION_REGISTRY.aggregate_functions() { + let name = func.name().to_string(); + session_state + .register_udaf(Arc::new(func)) + .context(RegisterUdfSnafu { name })?; + } + + let _ = session_state.register_udaf(quantile_udaf()); + + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); + let double_exponential_smoothing_udf = + DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); + let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); + + for udf in [ + NativeHistogramAbsentOverTime::scalar_udf(), + NativeHistogramAdd::scalar_udf(), + NativeHistogramAvg::scalar_udf(), + NativeHistogramAvgOverTime::scalar_udf(), + NativeHistogramChanges::scalar_udf(), + NativeHistogramCount::scalar_udf(), + NativeHistogramCountOverTime::scalar_udf(), + NativeHistogramDelta::scalar_udf(), + NativeHistogramDivScalar::scalar_udf(), + NativeHistogramDrop::bool_false_udf(String::new(), None), + NativeHistogramDrop::bool_true_udf(String::new(), None), + NativeHistogramDrop::float_null_udf(String::new(), None), + NativeHistogramEq::scalar_udf(), + NativeHistogramFraction::scalar_udf(), + NativeHistogramIDelta::scalar_udf(), + NativeHistogramIRate::scalar_udf(), + NativeHistogramIncrease::scalar_udf(), + NativeHistogramLastOverTime::scalar_udf(), + MixedRange::float_udf(None), + MixedRange::histogram_udf(None), + NativeHistogramMulScalar::scalar_udf(), + NativeHistogramNeg::scalar_udf(), + NativeHistogramNotEq::scalar_udf(), + NativeHistogramPresentOverTime::scalar_udf(), + NativeHistogramQuantile::scalar_udf(), + NativeHistogramRate::scalar_udf(), + NativeHistogramResets::scalar_udf(), + NativeHistogramScalarMul::scalar_udf(), + NativeHistogramStddev::scalar_udf(), + NativeHistogramStdvar::scalar_udf(), + NativeHistogramSub::scalar_udf(), + NativeHistogramSum::scalar_udf(), + NativeHistogramSumOverTime::scalar_udf(), + NativeHistogramToString::scalar_udf(), + PromqlFloatToString::scalar_udf(), + ] { + let _ = session_state.register_udf(Arc::new(udf)); + } + + for udaf in [ + NativeHistogramAggAvg::aggregate_udf(), + NativeHistogramAggSum::aggregate_udf(), + ] { + let _ = session_state.register_udaf(Arc::new(udaf)); + } + + Ok(()) +} + #[async_trait::async_trait] impl SubstraitPlanDecoder for DefaultPlanDecoder { async fn decode( @@ -123,108 +422,14 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_list: Arc, optimize: bool, ) -> common_query::error::Result { - // The session_state already has the `DefaultSerialzier` as `SerializerRegistry`. let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) + .with_serializer_registry(Arc::new(DefaultSerializer)) .with_catalog_list(catalog_list) .build(); - // Substrait decoder will look up the UDFs in SessionState, so we need to register them - // Note: the query context must be passed to set the timezone - // We MUST register the UDFs after we build the session state, otherwise the UDFs will be lost - // if they have the same name as the default UDFs or their alias. - // e.g. The default UDF `to_char()` has an alias `date_format()`, if we register a UDF with the name `date_format()` - // before we build the session state, the UDF will be lost. - for func in FUNCTION_REGISTRY.scalar_functions() { - let udf = func.provide(FunctionContext { - query_ctx: self.query_ctx.clone(), - state: Default::default(), - }); - session_state - .register_udf(Arc::new(udf)) - .context(RegisterUdfSnafu { name: func.name() })?; - } + // Re-register after the build to avoid Greptime UDF alias collisions. + register_greptime_functions(&mut session_state, &self.query_ctx)?; - for func in FUNCTION_REGISTRY.aggregate_functions() { - let name = func.name().to_string(); - session_state - .register_udaf(Arc::new(func)) - .context(RegisterUdfSnafu { name })?; - } - - let _ = session_state.register_udaf(quantile_udaf()); - - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); - let double_exponential_smoothing_udf = - DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); - let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); - - for udf in [ - NativeHistogramAbsentOverTime::scalar_udf(), - NativeHistogramAdd::scalar_udf(), - NativeHistogramAvg::scalar_udf(), - NativeHistogramAvgOverTime::scalar_udf(), - NativeHistogramChanges::scalar_udf(), - NativeHistogramCount::scalar_udf(), - NativeHistogramCountOverTime::scalar_udf(), - NativeHistogramDelta::scalar_udf(), - NativeHistogramDivScalar::scalar_udf(), - NativeHistogramDrop::bool_false_udf(String::new(), None), - NativeHistogramDrop::bool_true_udf(String::new(), None), - NativeHistogramDrop::float_null_udf(String::new(), None), - NativeHistogramEq::scalar_udf(), - NativeHistogramFraction::scalar_udf(), - NativeHistogramIDelta::scalar_udf(), - NativeHistogramIRate::scalar_udf(), - NativeHistogramIncrease::scalar_udf(), - NativeHistogramLastOverTime::scalar_udf(), - MixedRange::float_udf(None), - MixedRange::histogram_udf(None), - NativeHistogramMulScalar::scalar_udf(), - NativeHistogramNeg::scalar_udf(), - NativeHistogramNotEq::scalar_udf(), - NativeHistogramPresentOverTime::scalar_udf(), - NativeHistogramQuantile::scalar_udf(), - NativeHistogramRate::scalar_udf(), - NativeHistogramResets::scalar_udf(), - NativeHistogramScalarMul::scalar_udf(), - NativeHistogramStddev::scalar_udf(), - NativeHistogramStdvar::scalar_udf(), - NativeHistogramSub::scalar_udf(), - NativeHistogramSum::scalar_udf(), - NativeHistogramSumOverTime::scalar_udf(), - NativeHistogramToString::scalar_udf(), - PromqlFloatToString::scalar_udf(), - ] { - let _ = session_state.register_udf(Arc::new(udf)); - } - for udaf in [ - NativeHistogramAggAvg::aggregate_udf(), - NativeHistogramAggSum::aggregate_udf(), - ] { - let _ = session_state.register_udaf(Arc::new(udaf)); - } - - let logical_plan = DFLogicalSubstraitConvertor - .decode(message, session_state) + let logical_plan = decode_plan(message, &session_state, false) .await .map_err(BoxedError::new) .context(common_query::error::DecodePlanSnafu)?; @@ -241,18 +446,35 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { #[cfg(test)] mod tests { + use std::any::Any; + use std::sync::atomic::{AtomicBool, Ordering}; + + use catalog::error::Result as CatalogResult; + use catalog::{CatalogManager, RegisterTableRequest}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; + use common_function::aggrs::aggr_wrapper::StateWrapper; + use common_function::aggrs::aggr_wrapper::fix_order::FixStateUdafOrderingAnalyzer; + use common_query::logical_plan::SubstraitPlanDecoderRef; use common_query::native_histogram::native_histogram_value_type; + use common_time::Timezone; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; use datafusion::logical_expr::Extension; - use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, col, lit}; + use datafusion::optimizer::AnalyzerRule; + use datafusion_expr::expr::{AggregateFunction, Cast, InSubquery, ScalarFunction, Sort}; + use datafusion_expr::{ + AggregateUDF, Expr, LogicalPlanBuilder, LogicalTableSource, ScalarUDF, Subquery, col, lit, + }; use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, }; use datatypes::data_type::DataType; - use promql::extension_plan::RangeManipulate; + use futures::stream::BoxStream; + use promql::extension_plan::{RangeManipulate, SeriesNormalize, UnionDistinctOn}; use session::context::QueryContext; + use table::TableRef; + use table::metadata::{TableId, TableInfoRef}; + use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; use super::*; use crate::QueryEngineFactory; @@ -490,4 +712,946 @@ mod tests { assert!(decoded.contains("prom_mixed_range_float")); assert!(decoded.contains("prom_mixed_range_histogram")); } + + #[tokio::test] + async fn test_serializer_decode_merge_scan() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let plan = merge_scan(numbers_scan(), true); + let decoded = plan_decoder(engine.as_ref()) + .decode(encode_plan(&plan), request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + assert!(merge_scan.is_placeholder()); + assert_eq!(merge_scan.input().to_string(), numbers_scan().to_string()); + // `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping. + assert!(merge_scan.partition_cols().is_empty()); + } + + /// Payload table binding uses the engine catalog, not the request catalog. + #[tokio::test] + async fn test_serializer_decode_merge_scan_with_engine_catalog() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let decoded = plan_decoder(engine.as_ref()) + .decode( + encode_plan(&merge_scan(numbers_scan(), false)), + request_catalog_list(), + false, + ) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + let LogicalPlan::TableScan(scan) = merge_scan.input() else { + panic!("Expect a table scan, got: {}", merge_scan.input()); + }; + assert_eq!(scan.table_name.table(), NUMBERS_TABLE_NAME); + assert_eq!(scan.source.schema().fields().len(), 1); + assert_eq!(scan.source.schema().field(0).name(), "number"); + } + + /// Nested `MergeScan` payloads decode recursively. + #[tokio::test] + async fn test_serializer_decode_nested_merge_scan() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let inner = merge_scan(numbers_scan(), false); + let input = LogicalPlanBuilder::from(inner) + .filter(col("number").lt(lit(10u32))) + .unwrap() + .build() + .unwrap(); + let plan = merge_scan(input, false); + + let decoded = plan_decoder(engine.as_ref()) + .decode(encode_plan(&plan), request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + } + + fn greptime_date_format_udf(query_ctx: &QueryContextRef) -> Arc { + FUNCTION_REGISTRY + .get_function("date_format") + .expect("`date_format` is a function of GreptimeDB") + .provide(FunctionContext { + query_ctx: query_ctx.clone(), + state: Default::default(), + }) + .into() + } + + /// Uses a non-default timezone to distinguish the Greptime implementation. + fn date_format_query_ctx() -> QueryContextRef { + let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); + Arc::new(query_ctx) + } + + fn greptime_date_format_expr(query_ctx: &QueryContextRef, timestamp: Expr) -> Expr { + Expr::ScalarFunction(ScalarFunction { + func: greptime_date_format_udf(query_ctx), + args: vec![timestamp, lit("%Y-%m-%d %H:%i:%S")], + }) + } + + /// Includes `MergeScan` payload expressions, which are not node inputs. + fn scalar_functions_of(plan: &LogicalPlan) -> Vec> { + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + + let mut udfs = Vec::new(); + for expr in plan.expressions() { + expr.apply(|expr| { + if let Expr::ScalarFunction(func) = expr { + udfs.push(Arc::clone(&func.func)); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + } + + if let LogicalPlan::Extension(extension) = plan + && let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + { + udfs.extend(scalar_functions_of(merge_scan.input())); + } + + for input in plan.inputs() { + udfs.extend(scalar_functions_of(input)); + } + + udfs + } + + /// Asserts the Greptime `date_format` binding rather than DataFusion's `to_char` alias. + fn assert_only_function_is_greptime_date_format( + plan: &LogicalPlan, + query_ctx: &QueryContextRef, + ) { + let udfs = scalar_functions_of(plan); + assert_eq!(udfs.len(), 1, "expected one function in: {plan}"); + + let udf = &udfs[0]; + assert_eq!(udf.name(), "date_format", "got: {plan}"); + assert_eq!( + udf.as_ref(), + greptime_date_format_udf(query_ctx).as_ref(), + "`date_format` is not bound to the implementation of GreptimeDB: {plan}" + ); + } + + /// Request state registration keeps Greptime `date_format` ahead of DataFusion aliases. + #[tokio::test] + async fn test_serializer_decode_keeps_greptime_date_format() { + // Repeat to expose nondeterministic function-registration order. + for _ in 0..8 { + let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); + let factory = QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + + let table_provider = Arc::new(mock_table_provider(1.into())); + let query_ctx = date_format_query_ctx(); + let plan = LogicalPlanBuilder::scan( + "devices", + Arc::new(LogicalTableSource::new(table_provider.schema().clone())), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))]) + .unwrap() + .build() + .unwrap(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .unwrap() + .new_plan_decoder() + .unwrap(); + let decoded = plan_decoder + .decode( + bytes, + Arc::new(DummyCatalogList::with_table_provider(table_provider)), + false, + ) + .await + .unwrap(); + + assert_only_function_is_greptime_date_format(&decoded, &query_ctx); + } + } + + /// Nested payload decoding retains the Greptime `date_format` binding. + #[tokio::test] + async fn test_serializer_decode_merge_scan_payload_keeps_greptime_date_format() { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + + for _ in 0..8 { + let factory = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + let query_ctx = date_format_query_ctx(); + + let input = LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr( + &query_ctx, + Expr::Cast(Cast::new( + Box::new(col("number")), + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + )), + )]) + .unwrap() + .build() + .unwrap(); + let plan = + MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .unwrap() + .new_plan_decoder() + .unwrap(); + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))); + + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + + assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); + } + } + + /// Registers a `numbers` table in a memory catalog manager. + fn numbers_catalog_manager() -> CatalogManagerRef { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + catalog_manager + } + + /// Scan plan over the `numbers` schema shared by the fixtures below. + fn numbers_scan() -> LogicalPlan { + LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .build() + .unwrap() + } + + fn merge_scan(input: LogicalPlan, is_placeholder: bool) -> LogicalPlan { + MergeScanLogicalPlan::new(input, is_placeholder, Default::default()).into_logical_plan() + } + + fn encode_plan(plan: &LogicalPlan) -> Bytes { + DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .unwrap() + } + + fn plan_decoder(engine: &dyn crate::QueryEngine) -> SubstraitPlanDecoderRef { + engine + .engine_context(QueryContext::arc()) + .unwrap() + .new_plan_decoder() + .unwrap() + } + + /// Request catalog that resolves any table to a provider payload resolution must not use. + fn request_catalog_list() -> Arc { + Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))) + } + + /// Payload resolution preserves context and explicit catalog/schema components. + #[tokio::test] + async fn test_serializer_decode_payload_observes_catalog_and_query_context() { + let observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let inner = catalog::memory::new_memory_catalog_manager().unwrap(); + let manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager { + inner, + resolution: TableResolution::ObserveContext(observed.clone()), + }); + let mut resolver_ctx = QueryContext::arc(); + Arc::make_mut(&mut resolver_ctx).set_current_catalog("request_cat"); + Arc::make_mut(&mut resolver_ctx).set_current_schema("request_schema"); + let engine = QueryEngineFactory::new( + manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let mut state = engine + .engine_context(resolver_ctx.clone()) + .unwrap() + .state() + .clone(); + state.config_mut().set_extension(resolver_ctx); + let extensions = Extensions::default(); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager: state + .config() + .get_extension::() + .map(|s| s.catalog_manager().clone()), + payload: true, + }; + for reference in [ + TableReference::bare(NUMBERS_TABLE_NAME), + TableReference::partial("explicit_schema", NUMBERS_TABLE_NAME), + TableReference::full( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + NUMBERS_TABLE_NAME, + ), + ] { + consumer + .resolve_table_ref(&reference) + .await + .unwrap() + .unwrap(); + } + assert_eq!( + *observed.lock().unwrap(), + vec![ + [ + "request_cat".to_string(), + "request_schema".to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + [ + "request_cat".to_string(), + "explicit_schema".to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + [ + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + ], + ); + } + + /// How [`InterceptingCatalogManager`] resolves payload tables. + enum TableResolution { + /// Records catalog lookup and request context values. + ObserveContext(Arc>>), + /// Resolves once another task on the caller runtime signals. + AwaitSignal { + notify: Arc, + entered: Arc, + }, + } + + /// Engine catalog manager that takes over `table` resolution while delegating the rest. + struct InterceptingCatalogManager { + inner: CatalogManagerRef, + resolution: TableResolution, + } + + #[async_trait::async_trait] + impl CatalogManager for InterceptingCatalogManager { + fn as_any(&self) -> &dyn Any { + self.inner.as_any() + } + + async fn catalog_names(&self) -> CatalogResult> { + self.inner.catalog_names().await + } + + async fn schema_names( + &self, + catalog: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + self.inner.schema_names(catalog, query_ctx).await + } + + async fn table_names( + &self, + catalog: &str, + schema: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + self.inner.table_names(catalog, schema, query_ctx).await + } + + async fn catalog_exists(&self, catalog: &str) -> CatalogResult { + self.inner.catalog_exists(catalog).await + } + + async fn schema_exists( + &self, + catalog: &str, + schema: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult { + self.inner.schema_exists(catalog, schema, query_ctx).await + } + + async fn table_exists( + &self, + catalog: &str, + schema: &str, + table: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult { + self.inner + .table_exists(catalog, schema, table, query_ctx) + .await + } + + async fn table( + &self, + catalog: &str, + schema: &str, + table_name: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + match &self.resolution { + TableResolution::ObserveContext(observed) => { + let query_ctx = query_ctx.unwrap(); + observed.lock().unwrap().push([ + catalog.to_string(), + schema.to_string(), + table_name.to_string(), + query_ctx.current_catalog().to_string(), + query_ctx.current_schema(), + ]); + Ok(Some(NumbersTable::table(NUMBERS_TABLE_ID))) + } + + TableResolution::AwaitSignal { notify, entered } => { + entered.store(true, Ordering::SeqCst); + notify.notified().await; + self.inner + .table(catalog, schema, table_name, query_ctx) + .await + } + } + } + + async fn table_info_by_id(&self, table_id: TableId) -> CatalogResult> { + self.inner.table_info_by_id(table_id).await + } + + async fn tables_by_ids( + &self, + catalog: &str, + schema: &str, + table_ids: &[TableId], + ) -> CatalogResult> { + self.inner.tables_by_ids(catalog, schema, table_ids).await + } + + fn tables<'a>( + &'a self, + catalog: &'a str, + schema: &'a str, + query_ctx: Option<&'a QueryContext>, + ) -> BoxStream<'a, CatalogResult> { + self.inner.tables(catalog, schema, query_ctx) + } + } + + fn engine_with_intercepting_catalog( + resolution: TableResolution, + ) -> Arc { + let catalog_manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager { + inner: numbers_catalog_manager(), + resolution, + }); + QueryEngineFactory::new( + catalog_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + } + + /// A payload lookup parked on a caller-runtime task must keep the decode async. + /// + /// The engine catalog resolves only once a task spawned on this runtime signals it, so a + /// blocking bridge would deadlock the test instead of decoding. + #[tokio::test(flavor = "current_thread")] + async fn test_serializer_decode_merge_scan_awaits_caller_runtime_task() { + let notify = Arc::new(tokio::sync::Notify::new()); + let entered = Arc::new(AtomicBool::new(false)); + let engine = engine_with_intercepting_catalog(TableResolution::AwaitSignal { + notify: notify.clone(), + entered: entered.clone(), + }); + let plan = merge_scan(numbers_scan(), false); + let bytes = encode_plan(&plan); + + let decoder = plan_decoder(engine.as_ref()); + let mut decode = Box::pin(decoder.decode(bytes, request_catalog_list(), false)); + // First poll reaches the payload table lookup, which waits for the signal below. + assert!(futures::poll!(decode.as_mut()).is_pending()); + assert!(entered.load(Ordering::SeqCst)); + + let signal = tokio::spawn(async move { + notify.notify_waiters(); + }); + tokio::task::yield_now().await; + let decoded = decode.await.unwrap(); + signal.await.unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + } + + /// A missing engine-catalog table does not fall back to the usable request catalog. + #[tokio::test] + async fn test_serializer_decode_missing_engine_catalog_table() { + let missing_manager: CatalogManagerRef = + catalog::memory::new_memory_catalog_manager().unwrap(); + let engine = QueryEngineFactory::new( + missing_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let bytes = encode_plan(&merge_scan(numbers_scan(), false)); + + let err = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap_err(); + + let err = format!("{err:?}"); + assert!( + err.contains("Table not found: greptime.public.numbers"), + "unexpected error: {err}" + ); + } + + /// Manual states use their own catalog and query-context defaults for payloads. + #[tokio::test] + async fn test_serializer_decode_manual_state_keeps_extension_registry() { + let manager = catalog::memory::new_memory_catalog_manager().unwrap(); + manager.register_catalog_sync("manual_cat").unwrap(); + manager + .register_schema_sync(catalog::RegisterSchemaRequest { + catalog: "manual_cat".to_string(), + schema: "manual_schema".to_string(), + }) + .unwrap(); + manager + .register_table_sync(RegisterTableRequest { + catalog: "manual_cat".to_string(), + schema: "manual_schema".to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema")); + let mut config = datafusion::execution::context::SessionConfig::new() + .with_information_schema(true) + .with_default_catalog_and_schema("manual_cat", "manual_schema"); + config.set_extension(query_ctx.clone()); + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(config) + .build(); + let decoder = DefaultPlanDecoder::new(state, &query_ctx).unwrap(); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 0, + "number", + false, + Vec::new(), + merge_scan(numbers_scan(), false), + )), + }); + let request_catalog = Arc::new( + catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( + manager, query_ctx, + ), + ); + let decoded = decoder + .decode(encode_plan(&plan), request_catalog, false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow()); + } + + /// `MergeScan` decodes under an ordinary single-input extension parent. + #[tokio::test] + async fn test_serializer_decode_merge_scan_under_series_normalize() { + let outer = merge_scan(numbers_scan(), false); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 0, + "number", + false, + Vec::new(), + outer.clone(), + )), + }); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a SeriesNormalize plan, got: {decoded}"); + }; + let normalize = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a SeriesNormalize plan node"); + assert_eq!(normalize.inputs()[0].to_string(), outer.to_string()); + } + + /// `MergeScan` decodes under an ordinary multi-input extension parent. + #[tokio::test] + async fn test_serializer_decode_merge_scan_under_union_distinct_on() { + let left = merge_scan(numbers_scan(), false); + let right = merge_scan(numbers_scan(), true); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new( + UnionDistinctOn::try_new(left.clone(), right.clone(), vec![], 0).unwrap(), + ), + }); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a UnionDistinctOn plan, got: {decoded}"); + }; + let union = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a UnionDistinctOn plan node"); + assert_eq!(union.inputs()[0].to_string(), left.to_string()); + assert_eq!(union.inputs()[1].to_string(), right.to_string()); + } + + /// `MergeScan` inside a subquery still decodes through the wrapper consumer. + /// + /// Correlated outer references cannot be covered here: the in-repo Substrait producer + /// rejects `Expr::OuterReferenceColumn`, so no supported helper can build such a plan. + #[tokio::test] + async fn test_serializer_decode_merge_scan_inside_subquery() { + let inner = merge_scan(numbers_scan(), false); + let payload = LogicalPlanBuilder::from(numbers_scan()) + .filter(Expr::InSubquery(InSubquery::new( + Box::new(col("number")), + Subquery { + subquery: Arc::new(inner.clone()), + outer_ref_columns: vec![], + spans: Default::default(), + }, + false, + ))) + .unwrap() + .build() + .unwrap(); + let plan = merge_scan(payload.clone(), false); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + assert!( + decoded.to_string().contains("IN ()"), + "expect the decoded subquery filter, got: {decoded}" + ); + } + + /// Recursive payload decoding must restore ordered state aggregates, including their schema. + #[tokio::test] + async fn test_serializer_decode_payload_keeps_state_udaf_ordering() { + let last_value = (*datafusion::functions_aggregate::first_last::last_value_udaf()).clone(); + let state_udaf = AggregateUDF::new_from_impl(StateWrapper::new(last_value).unwrap()); + let aggr_expr = Expr::AggregateFunction(AggregateFunction::new_udf( + Arc::new(state_udaf), + vec![col("number")], + false, + None, + vec![Sort::new(col("number"), true, true)], + None, + )); + let payload = LogicalPlanBuilder::from(numbers_scan()) + .aggregate(Vec::::new(), vec![aggr_expr]) + .unwrap() + .build() + .unwrap(); + let bytes = encode_plan(&payload); + let encoded = Plan::decode(bytes.clone()).unwrap(); + let extensions = Extensions::try_from(&encoded.extensions).unwrap(); + let catalog_manager = numbers_catalog_manager(); + let engine = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let mut state = SessionStateBuilder::new_from_existing( + engine + .engine_context(QueryContext::arc()) + .unwrap() + .state() + .clone(), + ) + .with_serializer_registry(Arc::new(DefaultSerializer)) + .build(); + register_greptime_functions(&mut state, &QueryContext::arc()).unwrap(); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager: Some(catalog_manager), + payload: false, + }; + let decoded = consumer.decode_payload(bytes.to_vec()).await.unwrap(); + + let fixed_payload = FixStateUdafOrderingAnalyzer {} + .analyze(payload.clone(), &Default::default()) + .unwrap(); + assert_ne!( + fixed_payload.schema().as_arrow(), + payload.schema().as_arrow() + ); + assert_eq!( + decoded.schema().as_arrow(), + fixed_payload.schema().as_arrow() + ); + let LogicalPlan::Aggregate(aggregate) = &decoded else { + panic!("Expected an aggregate payload, got: {decoded}"); + }; + let Expr::AggregateFunction(function) = &aggregate.aggr_expr[0] else { + panic!("Expected a state aggregate, got: {:?}", aggregate.aggr_expr); + }; + assert!(function.func.inner().is::()); + assert_eq!(function.params.order_by.len(), 1); + let ArrowDataType::Struct(fields) = decoded.schema().field(0).data_type() else { + panic!("Expected the ordered aggregate state type"); + }; + assert_eq!(fields.len(), 3); + assert_eq!(fields[1].name(), "number"); + assert!(fields[1].is_nullable()); + } + + /// Type variation extensions keep being rejected while decoding. + #[tokio::test] + async fn test_serializer_decode_rejects_type_variations() { + use substrait::substrait_proto_df::proto::extensions::SimpleExtensionDeclaration; + use substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::{ + ExtensionTypeVariation, MappingType, + }; + + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let encoded = encode_plan(&merge_scan(numbers_scan(), false)); + let mut substrait_plan = Plan::decode(encoded).unwrap(); + substrait_plan.extensions.push(SimpleExtensionDeclaration { + mapping_type: Some(MappingType::ExtensionTypeVariation( + ExtensionTypeVariation { + extension_urn_reference: u32::MAX, + type_variation_anchor: 1, + name: "u!variation".to_string(), + }, + )), + }); + let bytes = Bytes::from(substrait_plan.encode_to_vec()); + + let err = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap_err(); + + let err = format!("{err:?}"); + assert!( + err.contains("Type variation extensions are not supported"), + "unexpected error: {err}" + ); + } } diff --git a/src/query/src/range_select/plan.rs b/src/query/src/range_select/plan.rs index 1961735c4af..d3a71196c9a 100644 --- a/src/query/src/range_select/plan.rs +++ b/src/query/src/range_select/plan.rs @@ -967,7 +967,9 @@ fn produce_align_time( // make modify_map for range_fn[i] for (row, hash) in by_columns_hash.iter().enumerate() { let ts = ts_column.value(row); - let ith_slot = (ts - align_to).div_floor(align); + let diff = ts - align_to; + // `div_euclid` equals `div_floor` for positive divisors (`align`). + let ith_slot = diff.div_euclid(align); let mut align_ts = ith_slot * align + align_to; while align_ts <= ts && ts < align_ts + range { modify_map diff --git a/src/query/src/sql/show_create_table.rs b/src/query/src/sql/show_create_table.rs index eb0392f15a1..c56e83d385a 100644 --- a/src/query/src/sql/show_create_table.rs +++ b/src/query/src/sql/show_create_table.rs @@ -23,11 +23,8 @@ use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, - COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, - COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, - COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, COMMENT_KEY, ColumnDefaultConstraint, ColumnSchema, - FulltextBackend, SchemaRef, + COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, + ColumnDefaultConstraint, ColumnSchema, FulltextBackend, SchemaRef, }; use datatypes::types::JsonFormat; use snafu::ResultExt; @@ -45,7 +42,7 @@ use table::requests::{ use crate::error::{ ConvertSqlTypeSnafu, ConvertSqlValueSnafu, GetFulltextOptionsSnafu, - GetSkippingIndexOptionsSnafu, GetVectorIndexOptionsSnafu, Result, SqlSnafu, + GetSkippingIndexOptionsSnafu, Result, SqlSnafu, }; /// Generates CREATE TABLE options from given table metadata and schema-level options. @@ -168,35 +165,6 @@ fn create_column(column_schema: &ColumnSchema, quote_style: char) -> Result, _>()?; + .collect::, _>>()?; let (min_val, max_val) = (typed_sorted_range[0].clone(), typed_sorted_range[1].clone()); @@ -2682,7 +2682,10 @@ mod test { let exec_stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); let real_output = exec_stream.collect::>().await; - let real_output: Vec<_> = real_output.into_iter().try_collect().unwrap(); + let real_output: Vec<_> = real_output + .into_iter() + .collect::, _>>() + .unwrap(); real_output } } diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index 006ac59771c..d20704cf616 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -41,3 +41,5 @@ cargo nextest run -p servers ``` Keep protocol translation here and permissions/database behavior in frontend. +Async table batching returns storage-completion futures with its early +acknowledgements. Keep completion reporting independent of the caller's lifetime. diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 0e2d896db8c..0be0e47eab7 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -58,7 +58,6 @@ common-telemetry.workspace = true common-time.workspace = true common-version = { workspace = true, features = ["codec"] } csv = "1.3" -dashmap.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true @@ -72,6 +71,7 @@ headers = "0.4" hex.workspace = true hostname = "0.3" http.workspace = true +http-body = "1" humantime.workspace = true humantime-serde.workspace = true hyper = { workspace = true, features = ["full"] } @@ -91,7 +91,7 @@ notify.workspace = true object-pool = "0.5" once_cell.workspace = true # Wait for https://github.com/databendlabs/opensrv/pull/81 -opensrv-mysql = { git = "https://github.com/GreptimeTeam/opensrv", rev = "f3135e5f00401e06075cf241c0cc26489e45b20a" } +opensrv-mysql = { git = "https://github.com/GreptimeTeam/opensrv", rev = "1f3d88ccaafb5721be6cfcc3ccbca48642426119" } opentelemetry-proto.workspace = true operator.workspace = true otel-arrow-rust.workspace = true @@ -123,7 +123,6 @@ serde.workspace = true serde_json.workspace = true session.workspace = true simd-json.workspace = true -simdutf8 = "0.1" smallvec.workspace = true snafu.workspace = true snap = "1" @@ -168,7 +167,7 @@ permutation = "0.4" rand.workspace = true serde_json.workspace = true servers = { workspace = true, features = ["testing"] } -session = { workspace = true, features = ["testing"] } +simdutf8 = "0.1" table.workspace = true tempfile = "3.0.0" tokio = { workspace = true, features = ["test-util"] } diff --git a/src/servers/benches/prom_decode.rs b/src/servers/benches/prom_decode.rs index 06b0d02e940..1cdc7d3782e 100644 --- a/src/servers/benches/prom_decode.rs +++ b/src/servers/benches/prom_decode.rs @@ -288,10 +288,9 @@ fn bench_prom_v1_vs_v2_decode(c: &mut Criterion) { group.bench_with_input(BenchmarkId::new("v2_manual", name), v2_bytes, |b, bytes| { b.iter(|| { black_box( - remote_write_v2::decode_uncompressed_write_requests( - black_box(bytes.as_slice()), - true, - ) + remote_write_v2::decode_uncompressed_write_requests(black_box( + bytes.as_slice(), + )) .unwrap(), ); }); diff --git a/src/servers/src/batcher.rs b/src/servers/src/batcher.rs index 10f8f558eb3..3cd9a831aaa 100644 --- a/src/servers/src/batcher.rs +++ b/src/servers/src/batcher.rs @@ -21,3 +21,60 @@ pub mod table; #[cfg(test)] mod test_util; + +use serde::{Deserialize, Serialize}; + +/// Controls whether batching waits for storage before replying to the client. +const PENDING_ROWS_BATCH_SYNC_ENV: &str = "PENDING_ROWS_BATCH_SYNC"; + +/// Returns whether pending-row batch submissions wait for the flush result +/// before replying to the client (synchronous mode), controlled by the +/// `PENDING_ROWS_BATCH_SYNC` environment variable and defaulting to `true`. +/// +/// Callers that reason about how long a remote write request may block (e.g. +/// the frontend HTTP timeout fallback) must consult this instead of +/// duplicating the env lookup. +pub fn pending_rows_batch_sync_enabled() -> bool { + std::env::var(PENDING_ROWS_BATCH_SYNC_ENV) + .ok() + .as_deref() + .and_then(|v| v.parse::().ok()) + .unwrap_or(true) +} + +/// Ingestion protocols that can opt into pending-row batching. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BatchingProtocol { + Prom, + Influxdb, + Opentsdb, + Otlp, + Logs, + Loki, + Splunk, + Elasticsearch, + HttpSql, + Mysql, + Postgres, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_protocol_names_reject_unknown_values() { + assert_eq!( + serde_json::from_str::("\"prom\"").unwrap(), + BatchingProtocol::Prom + ); + for name in ["sql", "unknown"] { + assert!(serde_json::from_str::(&format!("\"{name}\"")).is_err()); + } + assert_eq!( + serde_json::from_str::("\"http_sql\"").unwrap(), + BatchingProtocol::HttpSql + ); + } +} diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index f4c764ee8a8..7e0cea44b6a 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -21,6 +21,8 @@ mod tables; #[cfg(test)] mod test_util; +use std::collections::{HashMap, HashSet}; +use std::future::{Future, ready}; use std::num::NonZeroUsize; use std::sync::Arc; use std::time::Duration; @@ -34,11 +36,13 @@ use common_batcher::worker_registry::WorkerRegistry; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_time::timestamp::TimeUnit; use meter_core::data::MeterRecord; use meter_macros::write_meter; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::ResultExt; +use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; use tokio::sync::{Semaphore, broadcast, mpsc, oneshot}; use crate::batcher::flow_notifier::{FlowNotifier, start_flow_notification_worker}; @@ -54,6 +58,7 @@ pub use crate::batcher::logical_table::region_write::{ pub use crate::batcher::logical_table::tables::{ PendingRowsSchemaAlterer, PendingRowsSchemaAltererRef, }; +use crate::batcher::pending_rows_batch_sync_enabled; use crate::error; use crate::error::{Error, Result}; use crate::metrics::{ @@ -62,36 +67,18 @@ use crate::metrics::{ const PHYSICAL_TABLE_KEY: &str = "physical_table"; -/// Whether wait for ingestion result before reply to client. -const PENDING_ROWS_BATCH_SYNC_ENV: &str = "PENDING_ROWS_BATCH_SYNC"; - -/// Returns whether pending-row batch submissions wait for the flush result -/// before replying to the client (synchronous mode), controlled by the -/// `PENDING_ROWS_BATCH_SYNC` environment variable and defaulting to `true`. -/// -/// Callers that reason about how long a remote write request may block (e.g. -/// the frontend HTTP timeout fallback) must consult this instead of -/// duplicating the env lookup. -pub fn pending_rows_batch_sync_enabled() -> bool { - std::env::var(PENDING_ROWS_BATCH_SYNC_ENV) - .ok() - .as_deref() - .and_then(|v| v.parse::().ok()) - .unwrap_or(true) -} - const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; #[derive(Debug, Clone, Hash, Eq, PartialEq)] -struct BatchKey { - catalog: String, - schema: String, - physical_table: String, - skip_wal: bool, +pub(crate) struct BatchKey { + pub(crate) catalog: String, + pub(crate) schema: String, + pub(crate) physical_table: String, + pub(crate) skip_wal: bool, } // Requests can share a batch only when their write target and WAL policy match. -fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { +pub(crate) fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { let physical_table = ctx .extension(PHYSICAL_TABLE_KEY) .unwrap_or(GREPTIME_PHYSICAL_TABLE) @@ -182,14 +169,122 @@ impl LogicalTablePendingRowsBatcher { impl LogicalTablePendingRowsBatcher { pub async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result { + self.submit_with(requests, ctx, |_| ready(Ok(()))) + .await + .map(|(rows, ())| rows) + } + + /// Returns the physical metric table's time index unit resolved from + /// `ctx`, defaulting to millisecond when the table does not exist yet + /// (the schema alterer auto-creates it as millisecond). + async fn physical_time_index_unit_or_default(&self, ctx: &QueryContextRef) -> TimeUnit { + let key = batch_key_from_ctx(ctx); + let Ok(Some(table)) = self + .catalog_manager + .table(&key.catalog, &key.schema, &key.physical_table, None) + .await + else { + return TimeUnit::Millisecond; + }; + table + .table_info() + .meta + .schema + .timestamp_column() + .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit())) + .unwrap_or(TimeUnit::Millisecond) + } + + /// Returns whether the bulk path can accept `batches`: every existing + /// destination table must be a metric logical table bound to the + /// physical table selected by its context. A destination bound to + /// another physical table would be flushed through the selected + /// physical's regions, silently misplacing its rows, so such requests + /// must stay on the ordinary insert path (which routes per destination). + /// New tables are always fine: they are created on the selected physical + /// table. Time index units need no check here — the bulk encode converts + /// each request to its destination's unit. Destinations are resolved + /// once per distinct (schema, table). + pub(crate) async fn accepts_bulk_destinations( + &self, + batches: impl Iterator, + ) -> bool { + // One request can select different physical tables per batch (e.g. + // per-series physical-table labels), so the dedupe key includes the + // selected physical table: every distinct (schema, table, physical) + // triple is validated against the table's actual binding. + let mut checked = HashSet::new(); + // For missing tables, one request must not select two different + // physical tables: the batcher would create the table through one + // selection and flush its rows through the other's regions. + let mut missing_selections: HashMap<(String, String), String> = HashMap::new(); + for (ctx, requests) in batches { + let physical_table = batch_key_from_ctx(ctx).physical_table; + let schema = ctx.current_schema(); + for request in &requests.inserts { + if !checked.insert(( + schema.clone(), + request.table_name.clone(), + physical_table.clone(), + )) { + continue; + } + let Ok(Some(table)) = self + .catalog_manager + .table(ctx.current_catalog(), &schema, &request.table_name, None) + .await + else { + // New table: created on the selected physical table, but + // a conflicting selection within the same request cannot + // be batched. + if missing_selections + .insert( + (schema.clone(), request.table_name.clone()), + physical_table.clone(), + ) + .is_some_and(|previous| previous != physical_table) + { + return false; + } + continue; + }; + let info = table.table_info(); + if info.meta.engine != METRIC_ENGINE_NAME + || info + .meta + .options + .extra_options + .get(LOGICAL_TABLE_METADATA_KEY) + .map(String::as_str) + != Some(physical_table.as_str()) + { + return false; + } + } + } + true + } + + /// Submits with request-level accounting after schema preparation and before + /// queue admission. Acknowledgement follows the global batching policy. + pub async fn submit_with( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + after_prepare: impl FnOnce(RowInsertRequests) -> F, + ) -> Result<(u64, T)> + where + F: Future> + Send, + { let (table_batches, total_rows) = { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["submit_build_and_align"]) .start_timer(); - self.build_and_align_table_batches(requests, &ctx).await? + self.build_and_align_table_batches(&requests, &ctx).await? }; + let prepared = after_prepare(requests).await?; if total_rows == 0 { - return Ok(0); + return Ok((0, prepared)); } // Flushes dispatch directly to datanodes, so admit once before enqueueing. @@ -269,9 +364,9 @@ impl LogicalTablePendingRowsBatcher { }; result .context(error::SubmitBatchSnafu) - .map(|()| total_rows as u64) + .map(|()| (total_rows as u64, prepared)) } else { - Ok(total_rows as u64) + Ok((total_rows as u64, prepared)) } } } diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs index 28464eb2d3b..37080dbe258 100644 --- a/src/servers/src/batcher/logical_table/batch_convert.rs +++ b/src/servers/src/batcher/logical_table/batch_convert.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use arrow::compute::concat_batches; -use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; +use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use common_query::prelude::{greptime_timestamp, greptime_value}; use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse}; @@ -130,7 +130,7 @@ pub(in crate::batcher::logical_table) fn columns_taxonomy( essential_column_indices.push(index); } } - ArrowDataType::Timestamp(TimeUnit::Millisecond, _) => { + ArrowDataType::Timestamp(_, _) => { ensure!( timestamp_index.replace(index).is_none(), error::InvalidPromRemoteRequestSnafu { diff --git a/src/servers/src/batcher/logical_table/tables.rs b/src/servers/src/batcher/logical_table/tables.rs index 4ed07f73a6e..140dd2d4f23 100644 --- a/src/servers/src/batcher/logical_table/tables.rs +++ b/src/servers/src/batcher/logical_table/tables.rs @@ -15,9 +15,10 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use api::v1::{ColumnSchema, RowInsertRequests, Rows}; +use api::v1::{ColumnSchema, RowInsertRequests, Rows, SemanticType}; use arrow::datatypes::Schema as ArrowSchema; use async_trait::async_trait; +use common_time::timestamp::TimeUnit; use session::context::QueryContextRef; use snafu::OptionExt; @@ -75,7 +76,7 @@ impl LogicalTablePendingRowsBatcher { /// RecordBatches. pub(in crate::batcher::logical_table) async fn build_and_align_table_batches( &self, - requests: RowInsertRequests, + requests: &RowInsertRequests, ctx: &QueryContextRef, ) -> Result<(Vec<(String, u32, RecordBatchWithTsIdx)>, usize)> { let catalog = ctx.current_catalog().to_string(); @@ -91,6 +92,17 @@ impl LogicalTablePendingRowsBatcher { .plan_table_resolution(&catalog, &schema, ctx, &unique_tables) .await?; + // New tables are created on the request's selected physical table; + // their time index must use the physical table's unit (a missing + // physical table is auto-created as millisecond by the schema + // alterer, matching the default here). + if !plan.tables_to_create.is_empty() { + let physical_unit = self.physical_time_index_unit_or_default(ctx).await; + for (_, request_schema) in &mut plan.tables_to_create { + align_create_schema_time_index(request_schema, physical_unit); + } + } + self.create_missing_tables_and_refresh_schemas( &catalog, &schema, @@ -109,17 +121,28 @@ impl LogicalTablePendingRowsBatcher { } } +/// Rewrites the time index column of a create-table schema to `unit`, if the +/// schema carries a timestamp column in another unit. +fn align_create_schema_time_index(request_schema: &mut [ColumnSchema], unit: TimeUnit) { + for column in request_schema { + if column.semantic_type == SemanticType::Timestamp as i32 { + column.datatype = api::helper::timestamp_datatype(unit) as i32; + column.datatype_extension = None; + } + } +} + impl LogicalTablePendingRowsBatcher { /// Extracts non-empty `(table_name, rows)` pairs and computes total row /// count across the retained entries. pub(in crate::batcher::logical_table) fn collect_non_empty_table_rows( - requests: RowInsertRequests, - ) -> (Vec<(String, Rows)>, usize) { - let mut table_rows: Vec<(String, Rows)> = Vec::with_capacity(requests.inserts.len()); + requests: &RowInsertRequests, + ) -> (Vec<(&str, &Rows)>, usize) { + let mut table_rows: Vec<(&str, &Rows)> = Vec::with_capacity(requests.inserts.len()); let mut total_rows = 0; - for request in requests.inserts { - let Some(rows) = request.rows else { + for request in &requests.inserts { + let Some(rows) = &request.rows else { continue; }; if rows.rows.is_empty() { @@ -127,7 +150,7 @@ impl LogicalTablePendingRowsBatcher { } total_rows += rows.rows.len(); - table_rows.push((request.table_name, rows)); + table_rows.push((request.table_name.as_str(), rows)); } (table_rows, total_rows) @@ -137,15 +160,15 @@ impl LogicalTablePendingRowsBatcher { impl LogicalTablePendingRowsBatcher { /// Returns unique `(table_name, proto_schema)` pairs while keeping the /// first-seen schema for duplicate table names. - pub(in crate::batcher::logical_table) fn collect_unique_table_schemas( - table_rows: &[(String, Rows)], - ) -> Result> { + pub(in crate::batcher::logical_table) fn collect_unique_table_schemas<'a>( + table_rows: &[(&'a str, &'a Rows)], + ) -> Result> { let mut unique_tables: Vec<(&str, &[ColumnSchema])> = Vec::with_capacity(table_rows.len()); let mut seen = HashSet::new(); for (table_name, rows) in table_rows { - if seen.insert(table_name.as_str()) { - unique_tables.push((table_name.as_str(), &rows.schema)); + if seen.insert(*table_name) { + unique_tables.push((*table_name, &rows.schema)); } else { // table_rows should group rows by table name. return error::InvalidPromRemoteRequestSnafu { @@ -233,7 +256,7 @@ impl LogicalTablePendingRowsBatcher { catalog: &str, schema: &str, ctx: &QueryContextRef, - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], plan: &mut TableResolutionPlan, ) -> Result<()> { if plan.tables_to_create.is_empty() { @@ -298,7 +321,7 @@ impl LogicalTablePendingRowsBatcher { /// For newly created tables, re-checks all row schemas and appends alter /// operations when additional tag columns are still missing. pub(in crate::batcher::logical_table) fn enqueue_alter_for_new_tables( - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], plan: &mut TableResolutionPlan, ) -> Result<()> { let created_tables: HashSet<&str> = plan @@ -308,11 +331,11 @@ impl LogicalTablePendingRowsBatcher { .collect(); for (table_name, rows) in table_rows { - if !created_tables.contains(table_name.as_str()) { + if !created_tables.contains(table_name) { continue; } - let Some((region_schema, _)) = plan.region_schemas.get(table_name) else { + let Some((region_schema, _)) = plan.region_schemas.get(*table_name) else { continue; }; @@ -321,13 +344,13 @@ impl LogicalTablePendingRowsBatcher { || plan .tables_to_alter .iter() - .any(|(existing_name, _)| existing_name == table_name) + .any(|(existing_name, _)| existing_name == *table_name) { continue; } plan.tables_to_alter - .push((table_name.clone(), missing_columns)); + .push((table_name.to_string(), missing_columns)); } Ok(()) @@ -397,13 +420,13 @@ impl LogicalTablePendingRowsBatcher { /// Converts proto rows to `RecordBatch` values aligned to resolved region /// schemas and returns `(table_name, table_id, batch)` tuples. pub(in crate::batcher::logical_table) fn build_aligned_batches( - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], region_schemas: &HashMap, u32)>, ) -> Result> { let mut aligned_batches = Vec::with_capacity(table_rows.len()); for (table_name, rows) in table_rows { let (region_schema, table_id) = - region_schemas.get(table_name).cloned().with_context(|| { + region_schemas.get(*table_name).cloned().with_context(|| { error::UnexpectedResultSnafu { reason: format!("Region schema not resolved for table: {}", table_name), } @@ -415,7 +438,7 @@ impl LogicalTablePendingRowsBatcher { .start_timer(); rows_to_aligned_record_batch(rows, region_schema.as_ref())? }; - aligned_batches.push((table_name.clone(), table_id, record_batch)); + aligned_batches.push((table_name.to_string(), table_id, record_batch)); } Ok(aligned_batches) @@ -531,7 +554,7 @@ mod tests { }; let (table_rows, total_rows) = - LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(requests); + LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(&requests); assert_eq!(2, total_rows); assert_eq!(1, table_rows.len()); diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index 1a46a7f10f3..daae05251f8 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -30,6 +30,8 @@ use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; +use common_error::ext::BoxedError; +use common_query::{Output, WriteCompletion}; use operator::batcher::PendingRowsBatcher; use operator::error::{BatchFlushSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; @@ -38,6 +40,7 @@ use snafu::ResultExt; use table::metadata::TableInfoRef; use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, oneshot}; +use crate::batcher::pending_rows_batch_sync_enabled; use crate::batcher::table::flow_notifier::FlowNotifier; use crate::batcher::table::metrics::PENDING_WORKERS; use crate::batcher::table::pending_batch::PendingBatch; @@ -70,6 +73,7 @@ pub struct TablePendingRowsBatcher { flush_limiter: FlushLimiter, request_limiter: RequestLimiter, worker_channel_capacity: usize, + pending_rows_batch_sync: bool, worker_idle_timeout: Duration, inserter: Arc, flow_notifier: FlowNotifier, @@ -110,6 +114,7 @@ impl TablePendingRowsBatcher { flush_limiter, request_limiter, worker_channel_capacity, + pending_rows_batch_sync: pending_rows_batch_sync_enabled(), worker_idle_timeout: flush_interval.checked_mul(3).unwrap_or(flush_interval), inserter, flow_notifier, @@ -160,18 +165,18 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { }) } - /// Waits for completed bulk writes. Cancellation does not retract an admitted - /// submission. Combined failures affect all waiters and may be partial writes. + /// Acknowledges according to the global batching policy. Cancellation does + /// not retract an admitted submission. Flush failures may be partial writes. async fn submit( &self, table_info: TableInfoRef, batch: RecordBatch, ctx: QueryContextRef, permit: Arc, - ) -> Result { + ) -> Result { let total_rows = batch.num_rows(); if total_rows == 0 { - return Ok(0); + return Ok(Output::new_with_affected_rows(0)); } if table_info.catalog_name != ctx.current_catalog() || table_info.schema_name != ctx.current_schema() @@ -210,16 +215,30 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { } .fail(); } - response_rx - .await - .map_err(|_| { - UnexpectedSnafu { - violated: "batch worker stopped before reporting its write result".to_string(), - } - .build() - })? - .context(BatchFlushSnafu)?; - Ok(total_rows) + let completion = async move { + response_rx + .await + .map_err(|_| { + UnexpectedSnafu { + violated: "batch worker stopped before reporting its write result" + .to_string(), + } + .build() + })? + .context(BatchFlushSnafu) + }; + let mut output = Output::new_with_affected_rows(total_rows); + if self.pending_rows_batch_sync { + completion.await?; + } else { + output + .meta + .write_completions + .push(WriteCompletion::new(async move { + completion.await.map_err(BoxedError::new) + })); + } + Ok(output) } } @@ -236,7 +255,7 @@ mod tests { use api::v1::region::{RegionRequest, bulk_insert_request, region_request}; use api::v1::value::ValueData; use api::v1::{ColumnDataType, Row, Rows, Value}; - use arrow::array::{Int32Array, TimestampMillisecondArray}; + use arrow::array::{ArrayRef, Int32Array, TimestampMillisecondArray}; use arrow::datatypes::Schema as ArrowSchema; use arrow::record_batch::RecordBatch; use catalog::memory::MemoryCatalogManager; @@ -320,6 +339,120 @@ mod tests { } } + #[tokio::test] + async fn test_acknowledgement_preserves_admission() { + use operator::error::UnexpectedSnafu; + + use crate::batcher::pending_rows_batch_sync_enabled; + use crate::batcher::table::batch_key_from_ctx; + use crate::batcher::table::pending_batch::notify_batches; + + for sync in [false, true] { + for fail in [false, true] { + let backend = prepare_mocked_backend().await; + let nodes = Arc::new(MockDatanodeManager::new(BulkHandler { + requests: Arc::new(Mutex::new(Vec::new())), + report_missing_row: false, + expected_skip_wal: false, + expected_schema: "public".to_string(), + })); + let inserter = Arc::new(Inserter::new( + MemoryCatalogManager::new(), + create_partition_rule_manager(backend).await, + nodes, + mock_table_flownode_cache(1, vec![]).await, + true, + )); + let mut batcher = TablePendingRowsBatcher::try_new( + Duration::from_secs(3600), + 1, + 1, + 1, + 1, + NonZeroUsize::new(1).unwrap(), + inserter, + ) + .unwrap(); + assert_eq!( + batcher.pending_rows_batch_sync, + pending_rows_batch_sync_enabled() + ); + Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; + let table = Arc::new(new_test_table_info(1, "ack", [0].into_iter())); + let ctx = QueryContext::arc(); + let key = batch_key_from_ctx(&table.name, &ctx); + // Hold the worker command to control completion independently of scheduling. + let (_, receiver) = batcher.workers.get_or_create(key, 1).await; + let mut receiver = receiver.unwrap(); + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]) + .unwrap(); + let permit = batcher.acquire().await.unwrap(); + let submitter = batcher.clone(); + let submitted = + tokio::spawn(async move { submitter.submit(table, batch, ctx, permit).await }); + let WorkerCommand::Submit(pending) = + timeout(Duration::from_secs(5), receiver.recv()) + .await + .unwrap() + .unwrap(); + let mut submitted = Some(submitted); + let mut completion = None; + if sync { + assert!(!submitted.as_ref().unwrap().is_finished()); + } else { + let mut output = timeout(Duration::from_secs(5), submitted.take().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(output.extract_rows_and_cost().0, 1); + assert_eq!(output.meta.write_completions.len(), 1); + let mut pending = Box::pin(output.meta.write_completions.pop().unwrap().wait()); + assert!(futures::poll!(pending.as_mut()).is_pending()); + completion = Some(pending); + } + // Early acknowledgement must not release capacity before completion. + let acquire = batcher.acquire(); + tokio::pin!(acquire); + assert!(futures::poll!(acquire.as_mut()).is_pending()); + let result = if fail { + Err(Arc::new( + UnexpectedSnafu { + violated: "flush failed".to_string(), + } + .build(), + )) + } else { + Ok(()) + }; + notify_batches(vec![pending], result); + if let Some(submitted) = submitted { + let result = timeout(Duration::from_secs(5), submitted) + .await + .unwrap() + .unwrap(); + assert_eq!(result.is_err(), fail); + if !fail { + let output = result.unwrap(); + assert_eq!(output.extract_rows_and_cost().0, 1); + assert!(output.meta.write_completions.is_empty()); + } + } + if let Some(completion) = completion { + let result = timeout(Duration::from_secs(5), completion).await.unwrap(); + assert_eq!(result.is_err(), fail); + } + timeout(Duration::from_secs(5), acquire) + .await + .unwrap() + .unwrap(); + } + } + } + fn rows(value: i32) -> Rows { Rows { schema: vec![ @@ -422,9 +555,13 @@ mod tests { let (first_result, second_result) = timeout(Duration::from_secs(5), async { if shared_request { let permit = batcher.acquire().await.unwrap(); - tokio::join!( + let (first, second) = tokio::join!( batcher.submit(table.clone(), first, influx_ctx, permit.clone()), batcher.submit(table.clone(), second, opentsdb_ctx, permit), + ); + ( + first.map(|output| output.extract_rows_and_cost().0), + second.map(|output| output.extract_rows_and_cost().0), ) } else { let inserter = Inserter::new( diff --git a/src/servers/src/batcher/table/batch.rs b/src/servers/src/batcher/table/batch.rs index 21ebb94a33a..e9ab339a6ac 100644 --- a/src/servers/src/batcher/table/batch.rs +++ b/src/servers/src/batcher/table/batch.rs @@ -17,6 +17,7 @@ use std::time::Instant; use arrow::compute::concat_batches; use arrow::record_batch::RecordBatch; +use common_telemetry::error; use operator::error::{ComputeArrowSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; use operator::metrics::DIST_INGEST_ROW_COUNT; @@ -55,6 +56,7 @@ pub(in crate::batcher::table) async fn flush_batch( notifier.notify(table, &combined); } Err(error) => { + error!(error; "Failed to flush table batch, rows: {}", batch.total_rows); FLUSH_FAILURES.inc(); FLUSH_DROPPED_ROWS.inc_by(batch.total_rows as u64); notify_batches(batch.submissions, Err(Arc::new(error))); diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 47f9953a187..6e558c8234b 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -200,6 +200,9 @@ pub enum Error { location: Location, }, + #[snafu(display("Unsupported operation: {}", operation))] + UnsupportedOperation { operation: String }, + #[snafu(display("Invalid query: {}", reason))] InvalidQuery { reason: String, @@ -319,6 +322,18 @@ pub enum Error { error: std::io::Error, }, + #[snafu(display( + "Decompressed request body is too large: {} bytes exceeds the limit {} bytes", + size, + limit + ))] + DecompressedBodyTooLarge { + size: u64, + limit: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to compress prometheus remote request"))] CompressPromRemoteRequest { #[snafu(implicit)] @@ -755,7 +770,7 @@ impl ErrorExt for Error { | InvalidPromRemoteReadQueryResult { .. } | OtlpMetricModeIncompatible { .. } => StatusCode::IllegalState, - UnsupportedDataType { .. } => StatusCode::Unsupported, + UnsupportedDataType { .. } | UnsupportedOperation { .. } => StatusCode::Unsupported, #[cfg(not(windows))] UpdateJemallocMetrics { .. } => StatusCode::Internal, @@ -785,6 +800,7 @@ impl ErrorExt for Error { | DecompressSnappyPromRemoteRequest { .. } | DecompressSnappyLokiRequest { .. } | DecompressZstdPromRemoteRequest { .. } + | DecompressedBodyTooLarge { .. } | InvalidPromRemoteRequest { .. } | InvalidFlightTicket { .. } | InvalidPrepareStatement { .. } @@ -928,7 +944,7 @@ impl From for Error { } } -fn log_error_if_necessary(error: &Error) { +pub(crate) fn log_error_if_necessary(error: &Error) { if error.status_code().should_log_error() { error!(error; "Failed to handle HTTP request "); } else { diff --git a/src/servers/src/grpc.rs b/src/servers/src/grpc.rs index 50dd0b69c4a..9e9335b422d 100644 --- a/src/servers/src/grpc.rs +++ b/src/servers/src/grpc.rs @@ -36,6 +36,7 @@ use common_grpc::channel_manager::{ }; use common_telemetry::{error, info, warn}; use futures::FutureExt; +use http::{HeaderName, Method}; use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_server::ArrowMetricsServiceServer; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; @@ -48,8 +49,11 @@ use tonic::transport::ServerTlsConfig; use tonic::transport::server::TcpIncoming; use tonic::{Request, Response, Status}; use tonic_reflection::server::v1::{ServerReflection, ServerReflectionServer}; +use tower_http::cors::{AllowHeaders, CorsLayer}; use crate::error::{AlreadyStartedSnafu, InternalSnafu, Result, StartGrpcSnafu, TcpBindSnafu}; +use crate::grpc::memory_limit::MemoryLimiterExtensionService; +use crate::http::cors_allow_origin; use crate::install_default_crypto_provider; use crate::metrics::MetricsMiddlewareLayer; use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler}; @@ -86,6 +90,11 @@ pub struct GrpcOptions { /// The HTTP/2 keep-alive timeout. #[serde(with = "humantime_serde")] pub http2_keep_alive_timeout: Duration, + /// Whether to enable CORS, required by gRPC-Web clients in browsers. + /// Only the frontend's public gRPC server honors it. + pub enable_cors: bool, + /// Origins allowed by CORS. Empty allows any origin. + pub cors_allowed_origins: Vec, } impl GrpcOptions { @@ -154,6 +163,8 @@ impl Default for GrpcOptions { max_connection_age: None, http2_keep_alive_interval: Duration::from_secs(10), http2_keep_alive_timeout: Duration::from_secs(3), + enable_cors: false, + cors_allowed_origins: Vec::new(), } } } @@ -175,6 +186,8 @@ impl GrpcOptions { max_connection_age: None, http2_keep_alive_interval: Duration::from_secs(10), http2_keep_alive_timeout: Duration::from_secs(3), + enable_cors: false, + cors_allowed_origins: Vec::new(), } } @@ -213,6 +226,14 @@ impl FlightCompression { } } +/// The wrapped OTLP Arrow service type used by [`GrpcServer`]. +type OtelArrowService = MemoryLimiterExtensionService< + InterceptedService< + ArrowMetricsServiceServer>, + HeaderInterceptor, + >, +>; + pub struct GrpcServer { // states shutdown_tx: Mutex>>, @@ -224,17 +245,12 @@ pub struct GrpcServer { // tls config tls_config: Option, // Otel arrow service - otel_arrow_service: Mutex< - Option< - InterceptedService< - ArrowMetricsServiceServer>, - HeaderInterceptor, - >, - >, - >, + otel_arrow_service: Mutex>, bind_addr: Option, name: Option, config: GrpcServerConfig, + /// `None` disables CORS. + cors_allowed_origins: Option>, } /// Grpc Server configuration @@ -263,6 +279,24 @@ impl Default for GrpcServerConfig { } impl GrpcServer { + fn cors_layer(&self) -> Result> { + let Some(allowed_origins) = &self.cors_allowed_origins else { + return Ok(None); + }; + let layer = CorsLayer::new() + .allow_methods([Method::POST]) + .allow_origin(cors_allow_origin(allowed_origins)?) + .allow_headers(AllowHeaders::any()) + // Trailers-only responses (most errors) carry `grpc-status` in HTTP + // headers, which browsers hide unless exposed. + .expose_headers([ + HeaderName::from_static("grpc-status"), + HeaderName::from_static("grpc-message"), + HeaderName::from_static("grpc-status-details-bin"), + ]); + Ok(Some(layer)) + } + pub fn create_healthcheck_service(&self) -> HealthCheckServer { HealthCheckServer::new(HealthCheckHandler) } @@ -357,13 +391,18 @@ impl Server for GrpcServer { (incoming, addr) }; - let metrics_layer = tower::ServiceBuilder::new() + let cors_layer = self.cors_layer()?; + + // `Server::layer` prepends, so the first layer added is outermost. CORS must + // sit outside `GrpcWebLayer`, which rejects the `OPTIONS` preflight with a 400. + let middleware_layer = tower::ServiceBuilder::new() .layer(MetricsMiddlewareLayer) + .option_layer(cors_layer) .into_inner(); let mut builder = tonic::transport::Server::builder() .accept_http1(true) - .layer(metrics_layer) + .layer(middleware_layer) .layer(tonic_web::GrpcWebLayer::new()); if let Some(tls_config) = self.tls_config.clone() { @@ -423,9 +462,23 @@ impl Server for GrpcServer { #[cfg(test)] mod tests { - use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; + use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; - use super::{DEFAULT_GRPC_ADDR_PORT, format_server_addr, port_from_bind_addr}; + use common_runtime::Runtime; + use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; + use http::header::{ + ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN, + ACCESS_CONTROL_EXPOSE_HEADERS, ACCESS_CONTROL_REQUEST_HEADERS, + ACCESS_CONTROL_REQUEST_METHOD, CONTENT_TYPE, ORIGIN, + }; + use http::{HeaderName, Method, StatusCode}; + + use super::{ + DEFAULT_GRPC_ADDR_PORT, GrpcServer, GrpcServerConfig, format_server_addr, + port_from_bind_addr, + }; + use crate::grpc::builder::GrpcServerBuilder; + use crate::server::Server; #[test] fn test_port_from_bind_addr() { @@ -452,4 +505,118 @@ mod tests { format_server_addr(IpAddr::V6(Ipv6Addr::LOCALHOST), 3002) ); } + + /// The returned server owns the shutdown sender and must outlive the requests. + async fn start_test_grpc_server( + cors_allowed_origins: Option>, + ) -> (GrpcServer, SocketAddr) { + let runtime = Runtime::builder().build().unwrap(); + let mut builder = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime); + if let Some(origins) = cors_allowed_origins { + builder = builder.with_cors(origins); + } + let mut server = builder.build(); + server + .start(SocketAddr::from((Ipv4Addr::LOCALHOST, 0))) + .await + .unwrap(); + let addr = server.bind_addr().unwrap(); + (server, addr) + } + + const TEST_PATH: &str = "greptime.v1.HealthCheck/Check"; + + async fn send_preflight(addr: SocketAddr, origin: &str) -> reqwest::Response { + reqwest::Client::new() + .request(Method::OPTIONS, format!("http://{addr}/{TEST_PATH}")) + .header(ORIGIN, origin) + .header(ACCESS_CONTROL_REQUEST_METHOD, "POST") + .header(ACCESS_CONTROL_REQUEST_HEADERS, "content-type,x-grpc-web") + .send() + .await + .unwrap() + } + + /// The path is not served, so the response is a trailers-only `Unimplemented`. + async fn send_grpc_web_call(addr: SocketAddr, origin: &str) -> reqwest::Response { + reqwest::Client::new() + .post(format!("http://{addr}/{TEST_PATH}")) + .header(ORIGIN, origin) + .header(CONTENT_TYPE, "application/grpc-web+proto") + .header("x-grpc-web", "1") + // Empty gRPC-Web frame: 1-byte compression flag + 4-byte length. + .body(vec![0u8; 5]) + .send() + .await + .unwrap() + } + + fn header(response: &reqwest::Response, name: HeaderName) -> Option<&str> { + response.headers().get(name).map(|v| v.to_str().unwrap()) + } + + #[tokio::test] + async fn test_grpc_web_cors_preflight() { + let (_server, addr) = start_test_grpc_server(Some(Vec::new())).await; + + let response = send_preflight(addr, "https://example.com").await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), Some("*")); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_HEADERS), Some("*")); + // Lists the method of the actual call, not `OPTIONS`. + assert_eq!( + header(&response, ACCESS_CONTROL_ALLOW_METHODS), + Some("POST") + ); + } + + #[tokio::test] + async fn test_grpc_web_cors_expose_headers() { + let (_server, addr) = start_test_grpc_server(Some(Vec::new())).await; + + let response = send_grpc_web_call(addr, "https://example.com").await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), Some("*")); + // Premise: the response is trailers-only. + assert_eq!( + header(&response, HeaderName::from_static("grpc-status")), + Some("12") + ); + let expose_headers = header(&response, ACCESS_CONTROL_EXPOSE_HEADERS).unwrap(); + for name in ["grpc-status", "grpc-message", "grpc-status-details-bin"] { + assert!(expose_headers.contains(name), "{expose_headers}"); + } + } + + #[tokio::test] + async fn test_grpc_web_cors_custom_origins() { + let (_server, addr) = + start_test_grpc_server(Some(vec!["https://example.com".to_string()])).await; + + let response = send_preflight(addr, "https://example.com").await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), + Some("https://example.com") + ); + + // A disallowed origin still gets a 200, just without allow-origin; the + // browser does the blocking. + let response = send_preflight(addr, "https://notallowed.com").await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), None); + } + + #[tokio::test] + async fn test_grpc_web_cors_disabled() { + let (_server, addr) = start_test_grpc_server(None).await; + + let response = send_preflight(addr, "https://example.com").await; + + // Without CORS, `GrpcWebLayer` rejects the preflight. + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), None); + } } diff --git a/src/servers/src/grpc/builder.rs b/src/servers/src/grpc/builder.rs index d240cde3f1b..eaa7b628961 100644 --- a/src/servers/src/grpc/builder.rs +++ b/src/servers/src/grpc/builder.rs @@ -34,12 +34,13 @@ use tonic::codegen::Service; use tonic::service::RoutesBuilder; use tonic::service::interceptor::InterceptedService; use tonic::transport::{Identity, ServerTlsConfig}; -use tower::Layer; +use tower::{Layer, ServiceBuilder}; use crate::grpc::database::DatabaseService; use crate::grpc::flight::{FlightCraftRef, FlightCraftWrapper}; use crate::grpc::frontend_grpc_handler::FrontendGrpcHandler; use crate::grpc::greptime_handler::GreptimeRequestHandler; +use crate::grpc::memory_limit::{MemoryLimiterExtensionLayer, MemoryLimiterExtensionService}; use crate::grpc::prom_query_gateway::PrometheusGatewayService; use crate::grpc::region_server::{RegionServerHandlerRef, RegionServerRequestHandler}; use crate::grpc::{GrpcServer, GrpcServerConfig}; @@ -71,6 +72,7 @@ macro_rules! add_service { let service_with_limiter = $crate::tower::ServiceBuilder::new() .layer(MemoryLimiterExtensionLayer::new( $builder.memory_limiter().clone(), + max_recv_message_size, )) .service(service_builder); @@ -87,12 +89,15 @@ pub struct GrpcServerBuilder { routes_builder: RoutesBuilder, tls_config: Option, otel_arrow_service: Option< - InterceptedService< - ArrowMetricsServiceServer>, - HeaderInterceptor, + MemoryLimiterExtensionService< + InterceptedService< + ArrowMetricsServiceServer>, + HeaderInterceptor, + >, >, >, memory_limiter: ServerMemoryLimiter, + cors_allowed_origins: Option>, } impl GrpcServerBuilder { @@ -108,9 +113,16 @@ impl GrpcServerBuilder { tls_config: None, otel_arrow_service: None, memory_limiter, + cors_allowed_origins: None, } } + /// Enables CORS for gRPC-Web clients in browsers. An empty list allows any origin. + pub fn with_cors(mut self, allowed_origins: Vec) -> Self { + self.cors_allowed_origins = Some(allowed_origins); + self + } + /// Set a global memory limiter for all server protocols. pub fn with_memory_limiter(mut self, limiter: ServerMemoryLimiter) -> Self { self.memory_limiter = limiter; @@ -185,6 +197,13 @@ impl GrpcServerBuilder { .accept_compressed(CompressionEncoding::Zstd) .send_compressed(CompressionEncoding::Zstd); let svc = InterceptedService::new(server, HeaderInterceptor {}); + // Same pre-decode memory admission as `add_service!`. + let svc = ServiceBuilder::new() + .layer(MemoryLimiterExtensionLayer::new( + self.memory_limiter.clone(), + self.config.max_recv_message_size, + )) + .service(svc); self.otel_arrow_service = Some(svc); self } @@ -246,6 +265,7 @@ impl GrpcServerBuilder { bind_addr: None, name: self.name, config: self.config, + cors_allowed_origins: self.cors_allowed_origins, } } } diff --git a/src/servers/src/grpc/database.rs b/src/servers/src/grpc/database.rs index 2a63a982bef..237d760269f 100644 --- a/src/servers/src/grpc/database.rs +++ b/src/servers/src/grpc/database.rs @@ -25,6 +25,7 @@ use session::context::Channel; use tonic::{Request, Response, Status, Streaming}; use crate::grpc::greptime_handler::GreptimeRequestHandler; +use crate::grpc::memory_limit::PreDecodeMemoryReservation; use crate::grpc::{TonicResult, cancellation}; use crate::hint_headers; use crate::request_memory_limiter::ServerMemoryLimiter; @@ -57,7 +58,20 @@ impl GreptimeDatabase for DatabaseService { remote_addr, hints ); - let _guard = if let Some(limiter) = request.extensions().get::() { + // Retain the pre-decode reservation for the whole request: the + // extension holding the guard would be dropped when the request is + // consumed below, but the post-decode charge is skipped while it is + // active. + let _pre_reservation = request + .extensions() + .get::() + .cloned(); + let _guard = if _pre_reservation.is_some() { + // Compressed requests already reserved the worst-case decoded + // size before tonic decompressed the message; skip the exact + // post-decode charge to avoid double accounting. + None + } else if let Some(limiter) = request.extensions().get::() { let message_size = request.get_ref().encoded_len() as u64; Some(limiter.acquire(message_size).await?) } else { @@ -117,16 +131,31 @@ impl GreptimeDatabase for DatabaseService { ); let limiter = request.extensions().get::().cloned(); + // For compressed streams the whole stream's decoding memory was + // reserved before tonic started decompressing; messages are decoded + // one at a time, so the reservation covers each of them. The + // reservation is retained below for the whole stream: the extension + // holding the guard would otherwise be dropped when the request is + // consumed, while per-message charges stay skipped. + let reservation = request + .extensions() + .get::() + .cloned(); + let pre_reserved = reservation.is_some(); let handler = self.handler.clone(); let request_future = async move { let mut affected_rows = 0; + // Hold the pre-decode reservation until the stream is exhausted. + let _reservation = reservation; let mut stream = request.into_inner(); while let Some(request) = stream.next().await { let request = request?; - let _guard = if let Some(limiter_ref) = &limiter { + let _guard = if pre_reserved { + None + } else if let Some(limiter_ref) = &limiter { let message_size = request.encoded_len() as u64; Some(limiter_ref.acquire(message_size).await?) } else { diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index 71c98b5ddfd..a7514e889d6 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -46,6 +46,7 @@ use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested; use query::options::FlowQueryExtensions; use session::context::{Channel, QueryContextRef}; use snafu::{IntoError, OptionExt, ResultExt, ensure}; +use table::requests::AUTO_CREATE_TABLE_KEY; use table::table_name::TableName; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; @@ -255,8 +256,14 @@ impl FlightCraft for GreptimeRequestHandler { let limiter = extensions.get::().cloned(); - let query_ctx = + let mut query_ctx = context_auth::create_query_context_from_grpc_metadata(&headers, &extensions)?; + // Bulk streams use the same schema-on-write switch as row inserts. + for (key, value) in hint_headers::extract_hints(&headers) { + if key == AUTO_CREATE_TABLE_KEY { + Arc::make_mut(&mut query_ctx).set_extension(key, value); + } + } context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?; const MAX_PENDING_RESPONSES: usize = 32; diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index aebc2c5cc66..81dd4d9adf2 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -323,7 +323,8 @@ mod tests { use common_time::Timezone; use query::options::FLOW_SCHEDULED_TIME_MILLIS; use session::hints::{ - INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, + INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY, + REMOTE_QUERY_ID_EXTENSION_KEY, }; use snafu::ResultExt; use tonic::Code; @@ -419,6 +420,10 @@ mod tests { vec![ ("auto_create_table".to_string(), "true".to_string()), ("read_preference".to_string(), "leader".to_string()), + ( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + "follower".to_string(), + ), ( REMOTE_QUERY_ID_EXTENSION_KEY.to_string(), "spoofed".to_string(), @@ -447,6 +452,8 @@ mod tests { ReadPreference::Leader )); assert_eq!(query_context.extension("auto_create_table"), Some("true")); + assert_eq!(query_context.extension(READ_PREFERENCE_HINT), None); + assert_eq!(query_context.extension(READ_PREFERENCE_EXTENSION_KEY), None); assert_ne!(query_context.remote_query_id(), Some("spoofed")); assert!( query_context diff --git a/src/servers/src/grpc/memory_limit.rs b/src/servers/src/grpc/memory_limit.rs index 722a2c68c08..c9c62e88dd1 100644 --- a/src/servers/src/grpc/memory_limit.rs +++ b/src/servers/src/grpc/memory_limit.rs @@ -12,22 +12,69 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! Aggregate memory admission for gRPC services. +//! +//! Tonic decompresses and decodes a complete request message (up to +//! `max_recv_message_size`, 512 MiB by default) *before* the typed handler is +//! invoked, so a handler that charges the [`ServerMemoryLimiter`] only sees +//! the request after the memory is already allocated. For transport-compressed +//! messages (`grpc-encoding: gzip/zstd`) a few KiB on the wire can expand to +//! hundreds of MiB during that window, which would bypass any aggregate quota. +//! +//! [`MemoryLimiterExtensionLayer`] therefore reserves the configured maximum +//! decoded message size for every compressed request *before* the inner +//! (tonic) service runs, and holds the reservation for the whole request. +//! Handlers can detect the reservation via [`PreDecodeMemoryReservation`] in +//! the request extensions and skip their own post-decode charge to avoid +//! double accounting. +//! +//! The reservation is intentionally conservative (it upper-bounds the decoded +//! size before it is known); with the default unlimited limiter it is a no-op. + +use std::convert::Infallible; +use std::sync::Arc; use std::task::{Context, Poll}; +use axum::response::IntoResponse; +use common_memory_manager::MemoryGuard; use futures::future::BoxFuture; +use http::Request; +use tonic::Status; +use tonic::codegen::Service; use tonic::server::NamedService; -use tower::{Layer, Service}; +use tower::Layer; use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; + +/// The gRPC request compression selector header. +const GRPC_ENCODING_HEADER: &str = "grpc-encoding"; +/// The "no compression" encoding value. +const IDENTITY_ENCODING: &str = "identity"; + +/// Present in the request extensions when memory for the (compressed) +/// request's decoded message was reserved before tonic decompression. +/// +/// Handlers that charge the [`ServerMemoryLimiter`] after decoding should skip +/// that charge when this marker is present: the reservation already covers the +/// peak decoding memory and stays alive for the whole request. +#[derive(Clone)] +pub(crate) struct PreDecodeMemoryReservation { + _guard: Arc>, +} #[derive(Clone)] pub struct MemoryLimiterExtensionLayer { limiter: ServerMemoryLimiter, + max_decoding_message_size: usize, } impl MemoryLimiterExtensionLayer { - pub fn new(limiter: ServerMemoryLimiter) -> Self { - Self { limiter } + pub fn new(limiter: ServerMemoryLimiter, max_decoding_message_size: usize) -> Self { + Self { + limiter, + max_decoding_message_size, + } } } @@ -38,6 +85,7 @@ impl Layer for MemoryLimiterExtensionLayer { MemoryLimiterExtensionService { inner: service, limiter: self.limiter.clone(), + max_decoding_message_size: self.max_decoding_message_size, } } } @@ -46,27 +94,199 @@ impl Layer for MemoryLimiterExtensionLayer { pub struct MemoryLimiterExtensionService { inner: S, limiter: ServerMemoryLimiter, + max_decoding_message_size: usize, } impl NamedService for MemoryLimiterExtensionService { const NAME: &'static str = S::NAME; } -impl Service> for MemoryLimiterExtensionService +impl Service> for MemoryLimiterExtensionService where - S: Service>, + S: Service, Error = Infallible> + Clone + Send + 'static, + S::Response: axum::response::IntoResponse, S::Future: Send + 'static, + ReqBody: Send + 'static, { - type Response = S::Response; - type Error = S::Error; + type Response = axum::response::Response; + type Error = Infallible; type Future = BoxFuture<'static, Result>; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } - fn call(&mut self, mut req: http::Request) -> Self::Future { - req.extensions_mut().insert(self.limiter.clone()); - Box::pin(self.inner.call(req)) + fn call(&mut self, mut req: Request) -> Self::Future { + // Own a clone of the service so the returned future does not borrow + // `self` (tower's load-shedding pattern). + let mut this = self.clone(); + Box::pin(async move { + req.extensions_mut().insert(this.limiter.clone()); + + let compressed = req + .headers() + .get(GRPC_ENCODING_HEADER) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| !value.eq_ignore_ascii_case(IDENTITY_ENCODING)); + + if compressed { + // A compressed message can expand up to the decoding limit + // inside tonic, before any handler runs. Reserve that worst + // case against the aggregate quota so the decoding phase is + // admitted as well. No-op for an unlimited limiter. + let reservation = this.max_decoding_message_size as u64; + match this.limiter.acquire(reservation).await { + Ok(guard) => { + req.extensions_mut().insert(PreDecodeMemoryReservation { + _guard: Arc::new(guard), + }); + } + Err(e) => { + return Ok(Status::resource_exhausted(format!( + "request memory limit exceeded: {e}" + )) + .into_http::() + .into_response()); + } + } + } + + match this.inner.call(req).await { + Ok(response) => Ok(response.into_response()), + Err(e) => match e {}, + } + }) + } +} + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + + use axum::body::Body; + use axum::response::IntoResponse; + use common_memory_manager::OnExhaustedPolicy; + use futures_util::future::BoxFuture; + use http::{HeaderValue, StatusCode}; + use tower::ServiceExt; + + use super::*; + use crate::grpc::memory_limit::MemoryLimiterExtensionLayer; + + #[derive(Clone)] + struct EchoService; + + impl Service> for EchoService { + type Response = axum::response::Response; + type Error = Infallible; + type Future = BoxFuture<'static, Result>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: Request) -> Self::Future { + let saw_reservation = req + .extensions() + .get::() + .is_some(); + let saw_limiter = req.extensions().get::().is_some(); + Box::pin(async move { + Ok(( + StatusCode::OK, + format!("reservation={saw_reservation} limiter={saw_limiter}"), + ) + .into_response()) + }) + } + } + + impl NamedService for EchoService { + const NAME: &'static str = "test.Echo"; + } + + #[tokio::test] + async fn test_inserts_limiter_for_uncompressed_requests() { + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder().body(Body::empty()).unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!( + &body[..], + &b"reservation=false limiter=true"[..], + "uncompressed requests must not be pre-reserved" + ); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_reserves_max_message_size_for_compressed_requests() { + let max_size = 512 * 1024 * 1024; + let limiter = ServerMemoryLimiter::new(1024 * 1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), max_size).layer(EchoService); + + // The inner service observes the reservation; use a oneshot call and + // check the limiter was charged while the call is in flight via the + // response extensions observed by the echo service. + let req = Request::builder() + .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("gzip")) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!( + &body[..], + &b"reservation=true limiter=true"[..], + "compressed requests must carry a pre-decode reservation" + ); + // The reservation is released once the inner call (request) finishes. + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_rejects_compressed_request_when_quota_exhausted() { + // Quota smaller than the decoding limit: the reservation cannot fit. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder() + .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("zstd")) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + // gRPC errors are HTTP 200 with a grpc-status header. + assert_eq!( + res.headers() + .get("grpc-status") + .and_then(|v| v.to_str().ok()), + Some("8"), + "expected RESOURCE_EXHAUSTED grpc status" + ); + } + + #[tokio::test] + async fn test_identity_encoding_is_not_reserved() { + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder() + .header( + GRPC_ENCODING_HEADER, + HeaderValue::from_static(IDENTITY_ENCODING), + ) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], &b"reservation=false limiter=true"[..]); } } diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index a1cd71a147d..30f2ba6ec2e 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -53,6 +53,7 @@ use tower_http::trace::TraceLayer; use self::authorize::AuthState; use self::result::table_result::TableResponse; +use crate::batcher::BatchingProtocol; use crate::batcher::logical_table::LogicalTablePendingRowsBatcher; use crate::elasticsearch; use crate::error::{ @@ -111,6 +112,7 @@ mod timeout; pub mod utils; mod workload_scheduler; +pub use memory_limit::decoded_body_accounting_middleware; use result::HttpOutputWriter; pub(crate) use timeout::DynamicTimeoutLayer; @@ -189,22 +191,6 @@ pub(crate) enum HttpServerKind { Api, } -/// HTTP write protocols eligible for the shared pending-row batcher. -/// Prometheus uses this selector only when metric-engine storage is disabled. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum BatchingProtocol { - Prom, - Influxdb, - Opentsdb, - Otlp, - Logs, - Loki, - Splunk, - Elasticsearch, - HttpSql, -} - #[derive(Default)] pub struct HttpServer { router: StdMutex, @@ -215,6 +201,7 @@ pub struct HttpServer { // server configs options: HttpOptions, batching_protocols: Vec, + logical_batching_protocols: Vec, bind_addr: Option, /// What this server instance exposes. See [`HttpServerKind`]. kind: HttpServerKind, @@ -235,8 +222,11 @@ pub fn is_api_listener_path(path: &str) -> bool { is_namespace(path, HTTP_API_PREFIX_WITHOUT_TRAILING_SLASH) || is_namespace(path, "/dashboard") } +#[derive(Clone)] +struct LogicalBatchingProtocols(Vec); + /// Sets a local-only write selector after authentication creates the context. -async fn set_http_write_batching( +pub async fn set_http_write_batching( State(protocol): State, mut req: Request, next: Next, @@ -245,8 +235,13 @@ async fn set_http_write_batching( .extensions() .get::>>() .is_some_and(|protocols| protocols.contains(&protocol)); + let logical_enabled = req + .extensions() + .get::() + .is_some_and(|protocols| protocols.0.contains(&protocol)); if let Some(ctx) = req.extensions_mut().get_mut::() { ctx.set_batching_enabled(enabled); + ctx.set_logical_batching_enabled(logical_enabled); } next.run(req).await } @@ -650,6 +645,7 @@ pub struct DashboardState { pub struct HttpServerBuilder { options: HttpOptions, batching_protocols: Vec, + logical_batching_protocols: Vec, user_provider: Option, router: Router, memory_limiter: ServerMemoryLimiter, @@ -660,13 +656,20 @@ impl HttpServerBuilder { Self { options, batching_protocols: Vec::new(), + logical_batching_protocols: Vec::new(), user_provider: None, router: Router::new(), memory_limiter: ServerMemoryLimiter::default(), } } - /// Selects HTTP write protocols allowed to use the shared batcher. + /// Selects HTTP protocols allowed to use logical-table batching. + pub fn with_logical_batching_protocols(mut self, protocols: Vec) -> Self { + self.logical_batching_protocols = protocols; + self + } + + /// Selects HTTP protocols allowed to use ordinary-table batching. pub fn with_batching_protocols(mut self, protocols: Vec) -> Self { self.batching_protocols = protocols; self @@ -679,7 +682,12 @@ impl HttpServerBuilder { } pub fn with_sql_handler(self, sql_handler: ServerSqlQueryHandlerRef) -> Self { - let sql_router = HttpServer::route_sql(ApiState { sql_handler }); + let batch_body_limit = if self.options.body_limit.0 == 0 { + 4 * 1024 * 1024 + } else { + (self.options.body_limit.0 as usize).min(4 * 1024 * 1024) + }; + let sql_router = HttpServer::route_sql(ApiState { sql_handler }, batch_body_limit); Self { router: self @@ -714,7 +722,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/influxdb"), - HttpServer::route_influxdb(handler), + HttpServer::route_influxdb(handler, self.memory_limiter.clone()), ), ..self } @@ -726,7 +734,6 @@ impl HttpServerBuilder { pipeline_handler: Option, prom_store_with_metric_engine: bool, prom_validation_mode: PromValidationMode, - experimental_enable_prometheus_native_histogram: bool, pending_rows_batcher: Option>, ) -> Self { let state = PromStoreState { @@ -734,8 +741,8 @@ impl HttpServerBuilder { pipeline_handler, prom_store_with_metric_engine, prom_validation_mode, - experimental_enable_prometheus_native_histogram, pending_rows_batcher, + memory_limiter: self.memory_limiter.clone(), }; Self { @@ -761,16 +768,11 @@ impl HttpServerBuilder { self, handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, - experimental_enable_exponential_histogram: bool, ) -> Self { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/otlp"), - HttpServer::route_otlp( - handler, - with_metric_engine, - experimental_enable_exponential_histogram, - ), + HttpServer::route_otlp(handler, with_metric_engine, self.memory_limiter.clone()), ), ..self } @@ -804,23 +806,23 @@ impl HttpServerBuilder { let router = self.router.nest( &format!("/{HTTP_API_VERSION}"), - HttpServer::route_pipelines(log_state.clone()), + HttpServer::route_pipelines(log_state.clone(), self.memory_limiter.clone()), ); // deprecated since v0.11.0. Use `/logs` and `/pipelines` instead. let router = router.nest( &format!("/{HTTP_API_VERSION}/events"), #[allow(deprecated)] - HttpServer::route_log_deprecated(log_state.clone()), + HttpServer::route_log_deprecated(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( &format!("/{HTTP_API_VERSION}/loki"), - HttpServer::route_loki(log_state.clone()), + HttpServer::route_loki(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( &format!("/{HTTP_API_VERSION}/elasticsearch"), - HttpServer::route_elasticsearch(log_state.clone()), + HttpServer::route_elasticsearch(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( @@ -832,7 +834,7 @@ impl HttpServerBuilder { let router = router.nest( &format!("/{HTTP_API_VERSION}/splunk"), - HttpServer::route_splunk(log_state), + HttpServer::route_splunk(log_state, self.memory_limiter.clone()), ); Self { router, ..self } @@ -863,7 +865,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/dashboards"), - HttpServer::route_dashboard(handler), + HttpServer::route_dashboard(handler, self.memory_limiter.clone()), ), ..self } @@ -894,6 +896,7 @@ impl HttpServerBuilder { HttpServer { options: self.options, batching_protocols: self.batching_protocols.clone(), + logical_batching_protocols: self.logical_batching_protocols.clone(), user_provider: self.user_provider, shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -928,6 +931,7 @@ impl HttpServerBuilder { let internal = HttpServer { options: self.options, batching_protocols: self.batching_protocols.clone(), + logical_batching_protocols: self.logical_batching_protocols.clone(), user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router.clone()), @@ -946,6 +950,7 @@ impl HttpServerBuilder { Some(HttpServer { options: api_options, batching_protocols: self.batching_protocols, + logical_batching_protocols: self.logical_batching_protocols, user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -1057,20 +1062,7 @@ impl HttpServer { Method::DELETE, Method::HEAD, ]) - .allow_origin(if self.options.cors_allowed_origins.is_empty() { - AllowOrigin::from(Any) - } else { - AllowOrigin::from( - self.options - .cors_allowed_origins - .iter() - .map(|s| { - HeaderValue::from_str(s.as_str()) - .context(InvalidHeaderValueSnafu) - }) - .collect::>>()?, - ) - }) + .allow_origin(cors_allow_origin(&self.options.cors_allowed_origins)?) .allow_headers(Any), ) } else { @@ -1099,6 +1091,9 @@ impl HttpServer { authorize::check_http_auth, )) .layer(Extension(Arc::new(self.batching_protocols.clone()))) + .layer(Extension(LogicalBatchingProtocols( + self.logical_batching_protocols.clone(), + ))) .layer(middleware::from_fn(hints::extract_hints)) .layer(middleware::from_fn(client_ip::log_error_with_client_ip)) .layer(middleware::from_fn( @@ -1178,9 +1173,13 @@ impl HttpServer { .with_state(metrics_handler) } - fn route_loki(log_state: LogState) -> Router { + fn route_loki(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route("/api/v1/push", routing::post(loki::loki_ingest)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1192,7 +1191,7 @@ impl HttpServer { .with_state(log_state) } - fn route_splunk(log_state: LogState) -> Router { + fn route_splunk(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route( "/services/collector/health", @@ -1220,6 +1219,10 @@ impl HttpServer { "/services/collector/raw/1.0", routing::post(splunk::handle_raw), ) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1231,7 +1234,10 @@ impl HttpServer { .with_state(log_state) } - fn route_elasticsearch(log_state: LogState) -> Router { + fn route_elasticsearch( + log_state: LogState, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() // Return fake responsefor HEAD '/' request. .route( @@ -1303,6 +1309,10 @@ impl HttpServer { axum::Json(serde_json::json!({})), )), ) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer(ServiceBuilder::new().layer(RequestDecompressionLayer::new())) .layer(middleware::from_fn_with_state( BatchingProtocol::Elasticsearch, @@ -1312,7 +1322,10 @@ impl HttpServer { } #[deprecated(since = "0.11.0", note = "Use `route_pipelines()` instead.")] - fn route_log_deprecated(log_state: LogState) -> Router { + fn route_log_deprecated( + log_state: LogState, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() .route("/logs", routing::post(event::log_ingester)) .route( @@ -1328,6 +1341,10 @@ impl HttpServer { routing::delete(event::delete_pipeline), ) .route("/pipelines/dryrun", routing::post(event::pipeline_dryrun)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1339,7 +1356,7 @@ impl HttpServer { .with_state(log_state) } - fn route_pipelines(log_state: LogState) -> Router { + fn route_pipelines(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route("/ingest", routing::post(event::log_ingester)) .route( @@ -1359,6 +1376,10 @@ impl HttpServer { routing::delete(event::delete_pipeline), ) .route("/pipelines/_dryrun", routing::post(event::pipeline_dryrun)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1370,14 +1391,24 @@ impl HttpServer { .with_state(log_state) } - fn route_sql(api_state: ApiState) -> Router { + fn route_sql(api_state: ApiState, batch_body_limit: usize) -> Router { Router::new() .route( "/capabilities", - routing::get(|| async { - axum::Json(serde_json::json!({"metric_packed_import": 1})) + routing::get(|State(state): State| async move { + let mut capabilities = + serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1}); + if state.sql_handler.supports_metric_batch_ddl() { + capabilities["metric_batch_ddl"] = serde_json::json!(1); + } + axum::Json(capabilities) }), ) + .route( + "/ddl/logical-tables", + routing::post(handler::create_logical_tables) + .layer(DefaultBodyLimit::max(batch_body_limit)), + ) .route( "/sql", routing::get(handler::sql).post(handler::sql).layer( @@ -1457,10 +1488,17 @@ impl HttpServer { .with_state(state) } - fn route_influxdb(influxdb_handler: InfluxdbLineProtocolHandlerRef) -> Router { + fn route_influxdb( + influxdb_handler: InfluxdbLineProtocolHandlerRef, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() .route("/write", routing::post(influxdb_write_v1)) .route("/api/v2/write", routing::post(influxdb_write_v2)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1487,12 +1525,16 @@ impl HttpServer { fn route_otlp( otlp_handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, - experimental_enable_exponential_histogram: bool, + memory_limiter: ServerMemoryLimiter, ) -> Router { Router::new() .route("/v1/metrics", routing::post(otlp::metrics)) .route("/v1/traces", routing::post(otlp::traces)) .route("/v1/logs", routing::post(otlp::logs)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1503,7 +1545,6 @@ impl HttpServer { )) .with_state(OtlpState { with_metric_engine, - experimental_enable_exponential_histogram, handler: otlp_handler, }) } @@ -1534,13 +1575,20 @@ impl HttpServer { } #[cfg(feature = "dashboard")] - fn route_dashboard(handler: DashboardHandlerRef) -> Router { + fn route_dashboard( + handler: DashboardHandlerRef, + memory_limiter: ServerMemoryLimiter, + ) -> Router { use crate::http::dashboard::{add_dashboard, delete_dashboard, list_dashboards}; Router::new() .route("/", routing::get(list_dashboards)) .route("/{dashboard_name}", routing::post(add_dashboard)) .route("/{dashboard_name}", routing::delete(delete_dashboard)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1549,7 +1597,10 @@ impl HttpServer { } #[cfg(not(feature = "dashboard"))] - fn route_dashboard(handler: DashboardHandlerRef) -> Router { + fn route_dashboard( + handler: DashboardHandlerRef, + _memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new().with_state(DashboardState { handler }) } } @@ -1557,6 +1608,18 @@ impl HttpServer { pub const HTTP_SERVER: &str = "HTTP_SERVER"; pub const HTTP_API_SERVER: &str = "HTTP_API_SERVER"; +/// An empty list allows any origin. +pub(crate) fn cors_allow_origin(origins: &[String]) -> Result { + if origins.is_empty() { + return Ok(AllowOrigin::any()); + } + let origins = origins + .iter() + .map(|origin| HeaderValue::from_str(origin).context(InvalidHeaderValueSnafu)) + .collect::>>()?; + Ok(AllowOrigin::list(origins)) +} + #[async_trait] impl Server for HttpServer { async fn shutdown(&self) -> Result<()> { @@ -1885,6 +1948,15 @@ mod test { client.get("/v1/capabilities").send().await.status(), StatusCode::UNAUTHORIZED ); + assert_eq!( + client + .post("/v1/ddl/logical-tables") + .form(&[("sql", "SELECT 1")]) + .send() + .await + .status(), + StatusCode::UNAUTHORIZED + ); let response = client .get("/v1/capabilities") .header("Authorization", "Basic dXNlcjpwYXNzd29yZA==") @@ -1893,11 +1965,89 @@ mod test { assert_eq!(response.status(), StatusCode::OK); assert_eq!( response.json::().await, - serde_json::json!({"metric_packed_import": 1}) + serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1}) ); } } + #[tokio::test] + async fn logical_ddl_form_limits_and_unsupported_handler() { + use axum::body::{Body, to_bytes}; + use tower::ServiceExt; + + async fn post_form(app: Router, body: String) -> (StatusCode, serde_json::Value) { + let response = app + .oneshot( + Request::builder() + .method(Method::POST) + .uri("/v1/ddl/logical-tables") + .header("Content-Type", "application/x-www-form-urlencoded") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + (status, serde_json::from_slice(&body).unwrap()) + } + + for limit in [0, 64] { + let (tx, _rx) = mpsc::channel(1); + let server = HttpServerBuilder::new(HttpOptions { + body_limit: ReadableSize(limit), + ..Default::default() + }) + .with_sql_handler(Arc::new(DummyInstance { _tx: tx })) + .build(); + let app = server.build(server.make_app()).unwrap(); + let client = TestClient::new(app.clone()).await; + assert_eq!( + client.get("/v1/ddl/logical-tables").send().await.status(), + StatusCode::METHOD_NOT_ALLOWED + ); + let small = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", "CREATE TABLE t")]) + .send() + .await; + assert_eq!( + small.json::().await["code"], + common_error::status_code::StatusCode::Unsupported as u32 + ); + let (_, body) = post_form( + app.clone(), + if limit == 0 { + format!("sql={}", "%20".repeat(1024 * 1024)) + } else { + "sql=".to_string() + &" ".repeat(65) + }, + ) + .await; + let code = body["code"].clone(); + assert_eq!( + code, + if limit == 0 { + common_error::status_code::StatusCode::Unsupported as u32 + } else { + common_error::status_code::StatusCode::InvalidArguments as u32 + } + ); + for body in [ + format!("sql={}", "x".repeat(1024 * 1024 + 1)), + format!("sql={}", "%C3%A9".repeat(512 * 1024 + 1)), + format!("sql=CREATE+TABLE+t&padding={}", "x".repeat(4 * 1024 * 1024)), + ] { + let (status, body) = post_form(app.clone(), body).await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!( + body["code"], + common_error::status_code::StatusCode::InvalidArguments as u32 + ); + } + } + } + #[tokio::test] pub async fn test_extra_router_respects_api_surface() { // Downstream routes added through `with_extra_router()` (e.g. enterprise @@ -2450,26 +2600,50 @@ mod batching_tests { use common_query::Output; use session::context::QueryContextRef; + use crate::batcher::BatchingProtocol; use crate::error::Result as ServerResult; use crate::http::test_helpers::TestClient; - use crate::http::{BatchingProtocol, HttpOptions, HttpServerBuilder}; + use crate::http::{HttpOptions, HttpServerBuilder}; use crate::influxdb::InfluxdbRequest; use crate::opentsdb::codec::DataPoint; use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler}; - #[test] - fn test_protocol_names_reject_unknown_values() { - assert_eq!( - serde_json::from_str::("\"prom\"").unwrap(), - BatchingProtocol::Prom - ); - for name in ["sql", "unknown"] { - assert!(serde_json::from_str::(&format!("\"{name}\"")).is_err()); + #[tokio::test] + async fn test_batching_selectors_are_independent() { + use axum::routing::post; + use axum::{Extension, Json, Router, middleware}; + use session::context::{QueryContext, QueryContextBuilder}; + + use crate::http::{LogicalBatchingProtocols, set_http_write_batching}; + for table in [false, true] { + for logical in [false, true] { + let app = Router::new() + .route( + "/", + post(|Extension(ctx): Extension| async move { + Json([ctx.batching_enabled(), ctx.logical_batching_enabled()]) + }), + ) + .route_layer(middleware::from_fn_with_state( + BatchingProtocol::Otlp, + set_http_write_batching, + )) + .layer(Extension(QueryContextBuilder::default().build())) + .layer(Extension(Arc::new(if table { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }))) + .layer(Extension(LogicalBatchingProtocols(if logical { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }))); + let client = TestClient::new(app).await; + let actual: [bool; 2] = client.post("/").send().await.json().await; + assert_eq!(actual, [table, logical]); + } } - assert_eq!( - serde_json::from_str::("\"http_sql\"").unwrap(), - BatchingProtocol::HttpSql - ); } #[derive(Default)] diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index bd31ee35ace..8cbfc25ff07 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -44,7 +44,10 @@ use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use tokio::sync::{Notify, watch}; -use crate::error::{CollectRecordbatchSnafu, FailedToParseQuerySnafu, InvalidQuerySnafu, Result}; +use crate::error::{ + CollectRecordbatchSnafu, FailedToParseQuerySnafu, InvalidQuerySnafu, Result, + log_error_if_necessary, +}; use crate::http::header::collect_plan_metrics; use crate::http::result::arrow_result::ArrowResponse; use crate::http::result::csv_result::CsvResponse; @@ -229,6 +232,83 @@ pub async fn sql( resp.with_execution_time(start.elapsed().as_millis() as u64) } +/// Form parameters for the bounded logical-table DDL operation. +#[derive(Deserialize)] +pub struct LogicalTableDdl { + sql: Option, + db: Option, +} + +/// Creates one batch of Metric logical tables using the normal v1 response. +pub async fn create_logical_tables( + State(state): State, + Extension(mut query_ctx): Extension, + form: std::result::Result, FormRejection>, +) -> HttpResponse { + let start = Instant::now(); + let form = match form { + Ok(Form(form)) => form, + Err(_) => { + return HttpResponse::Error(ErrorResponse::from_error_message( + StatusCode::InvalidArguments, + "invalid or oversized form body".to_string(), + )); + } + }; + let Some(sql) = form.sql else { + return HttpResponse::Error(ErrorResponse::from_error_message( + StatusCode::InvalidArguments, + "sql parameter is required.".to_string(), + )); + }; + if sql.len() > crate::query_handler::sql::MAX_LOGICAL_TABLE_DDL_BYTES { + return HttpResponse::Error(ErrorResponse::from_error_message( + StatusCode::InvalidArguments, + "logical-table DDL exceeds 1 MiB".to_string(), + )); + } + if let Some(db) = form.db { + let (catalog, schema) = parse_catalog_and_schema_from_db_string(&db); + query_ctx.set_current_catalog(&catalog); + query_ctx.set_current_schema(&schema); + } + query_ctx.set_channel(Channel::HttpSql); + let query_ctx = Arc::new(query_ctx); + if let Some((status, msg)) = validate_schema(state.sql_handler.clone(), query_ctx.clone()).await + { + return HttpResponse::Error(ErrorResponse::from_error_message(status, msg)); + } + let outputs = match state + .sql_handler + .create_logical_tables(&sql, query_ctx) + .await + { + Ok(outputs) + if outputs + .iter() + .all(|output| matches!(output.data, OutputData::AffectedRows(0))) => + { + outputs.into_iter().map(Ok).collect() + } + other => { + let status = match other { + Err(error) => { + log_error_if_necessary(&error); + error.status_code() + } + Ok(_) => StatusCode::Internal, + }; + return HttpResponse::Error( + ErrorResponse::from_error_message(status, "logical-table batch failed".to_string()) + .with_execution_time(start.elapsed().as_millis() as u64), + ); + } + }; + GreptimedbV1Response::from_output(outputs) + .await + .with_execution_time(start.elapsed().as_millis() as u64) +} + /// Limits each statement's response before collecting batches and converting rows. fn limit_output_rows(output: Output, limit: usize) -> Result { let mut remaining = limit; @@ -994,13 +1074,16 @@ mod tests { use std::sync::atomic::AtomicUsize; use async_trait::async_trait; + use common_error::ext::{BoxedError, PlainError}; use common_query::OutputMeta; + use common_telemetry::tracing::instrument::WithSubscriber; use datafusion_expr::LogicalPlan; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{BinaryVector, UInt32Vector, VectorRef}; use futures::stream; use query::query_engine::DescribeResult; + use snafu::IntoError; use super::*; use crate::query_handler::sql::SqlQueryHandler; @@ -1015,6 +1098,12 @@ mod tests { std::mem::take(&mut *self.outputs.lock().unwrap()) } + async fn create_logical_tables(&self, _: &str, _: QueryContextRef) -> Result> { + std::mem::take(&mut *self.outputs.lock().unwrap()) + .into_iter() + .collect() + } + async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result { unimplemented!() } @@ -1068,6 +1157,61 @@ mod tests { .await } + #[tokio::test] + async fn test_logical_ddl_logs_diagnostics_without_exposing_sql_or_response_details() { + let sql = "CREATE TABLE logs (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='request-only-value')"; + for (status, level) in [ + (StatusCode::InvalidArguments, "WARN"), + (StatusCode::Internal, "ERROR"), + ] { + let logs = tempfile::NamedTempFile::new().unwrap(); + let writer = logs.reopen().unwrap(); + let subscriber = common_telemetry::tracing_subscriber::fmt() + .without_time() + .with_ansi(false) + .with_writer(move || writer.try_clone().unwrap()) + .finish(); + let error = crate::error::ExecuteQuerySnafu.into_error(BoxedError::new( + PlainError::new("sensitive-option-value".to_string(), status), + )); + let response = create_logical_tables( + State(ApiState { + sql_handler: Arc::new(TestSqlQueryHandler { + outputs: Mutex::new(vec![Err(error)]), + }), + }), + Extension(QueryContext::with_db_name(None)), + Ok(Form(LogicalTableDdl { + sql: Some(sql.to_string()), + db: None, + })), + ) + .with_subscriber(subscriber) + .await + .into_response(); + let diagnostics = std::fs::read_to_string(logs.path()).unwrap(); + assert!(diagnostics.contains(level), "{diagnostics}"); + assert!( + diagnostics.contains("Failed to execute query"), + "{diagnostics}" + ); + assert!( + diagnostics.contains("sensitive-option-value"), + "{diagnostics}" + ); + assert!(!diagnostics.contains(sql), "{diagnostics}"); + assert!(!diagnostics.contains("request-only-value"), "{diagnostics}"); + assert!(!format!("{:?}", response.headers()).contains("sensitive-option-value")); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(body["code"], status as u32); + assert_eq!(body["error"], "logical-table batch failed"); + assert!(body.get("output").is_none()); + } + } + fn number_batches() -> RecordBatches { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "number", diff --git a/src/servers/src/http/hints.rs b/src/servers/src/http/hints.rs index e3e96f94c3b..9d5ca035d0c 100644 --- a/src/servers/src/http/hints.rs +++ b/src/servers/src/http/hints.rs @@ -68,7 +68,8 @@ mod tests { use query::options::FLOW_SCHEDULED_TIME_MILLIS; use session::context::{QueryContextBuilder, generate_remote_query_id}; use session::hints::{ - INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, + INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY, + REMOTE_QUERY_ID_EXTENSION_KEY, }; use super::apply_hints; @@ -94,6 +95,10 @@ mod tests { INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(), "spoofed-regs".to_string(), ), + ( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + "follower".to_string(), + ), ( FLOW_SCHEDULED_TIME_MILLIS.to_string(), "1700000000000".to_string(), @@ -112,6 +117,8 @@ mod tests { .extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY) .is_none() ); + assert_eq!(query_ctx.read_preference(), session::ReadPreference::Leader); + assert!(query_ctx.extension(READ_PREFERENCE_EXTENSION_KEY).is_none()); assert_eq!( query_ctx.extension(FLOW_SCHEDULED_TIME_MILLIS), Some("1700000000000") diff --git a/src/servers/src/http/jaeger.rs b/src/servers/src/http/jaeger.rs index 32ef6a61682..16dbaf76e53 100644 --- a/src/servers/src/http/jaeger.rs +++ b/src/servers/src/http/jaeger.rs @@ -328,7 +328,7 @@ impl QueryTraceParams { } .build() })?; - for (_, v) in tags_map.iter_mut() { + for v in tags_map.values_mut() { if let Some(number) = convert_string_to_number(v) { *v = number; } diff --git a/src/servers/src/http/loki.rs b/src/servers/src/http/loki.rs index 4e12eeb931e..161664a2ce1 100644 --- a/src/servers/src/http/loki.rs +++ b/src/servers/src/http/loki.rs @@ -108,6 +108,7 @@ pub async fn loki_ingest( TypedHeader(content_type): TypedHeader, LogTableName(table_name): LogTableName, pipeline_info: PipelineInfo, + Extension(memory_limiter): Extension, bytes: Bytes, ) -> Result { ctx.set_channel(Channel::Loki); @@ -127,6 +128,7 @@ pub async fn loki_ingest( pipeline_info, bytes, &ctx, + &memory_limiter, ) .await?; @@ -182,6 +184,7 @@ async fn build_loki_context_req( pipeline_info: PipelineInfo, bytes: Bytes, ctx: &QueryContextRef, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { // A pipeline header switches Loki into the generic pipeline path; without // it, Loki writes directly to the target log table. @@ -195,10 +198,12 @@ async fn build_loki_context_req( pipeline_params: pipeline_info.pipeline_params, bytes, }; - build_loki_pipeline_context_req(handler, pipeline_req, ctx).await + build_loki_pipeline_context_req(handler, pipeline_req, ctx, memory_limiter).await } None => { - let req = build_loki_raw_insert_request(content_type, table_name, bytes)?; + let req = + build_loki_raw_insert_request(content_type, table_name, bytes, memory_limiter) + .await?; Ok(ContextReq::default_opt_with_reqs(vec![req])) } } @@ -208,6 +213,7 @@ async fn build_loki_pipeline_context_req( handler: &PipelineHandlerRef, pipeline_req: LokiPipelineContextReq, ctx: &QueryContextRef, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { let LokiPipelineContextReq { content_type, @@ -223,7 +229,8 @@ async fn build_loki_pipeline_context_req( PipelineDefinition::from_name(&pipeline_name, version, None).context(PipelineSnafu)?; let pipeline_ctx = PipelineContext::new(&def, &pipeline_params, Channel::Loki); - let values = extract_item::(content_type, bytes)? + let values = extract_item::(content_type, bytes, memory_limiter) + .await? .map(|item| item.map) .collect::>(); @@ -235,14 +242,16 @@ async fn build_loki_pipeline_context_req( run_pipeline(handler, &pipeline_ctx, req, ctx, true).await } -fn build_loki_raw_insert_request( +async fn build_loki_raw_insert_request( content_type: ContentType, table_name: String, bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { let mut schema_info = SchemaInfo::from_schema_list(LOKI_INIT_SCHEMAS.clone()); let mut rows = Vec::with_capacity(256); - for loki_row in extract_item::(content_type, bytes)? { + let loki_rows = extract_item::(content_type, bytes, memory_limiter).await?; + for loki_row in loki_rows { let mut row = init_row( schema_info.schema.len(), loki_row.ts, @@ -275,7 +284,11 @@ fn build_loki_raw_insert_request( /// JSON push bodies become `LokiMiddleItem`, protobuf push bodies /// become `LokiMiddleItem>`, and the generic `Into` /// conversion selects either direct-write `LokiRawItem` or pipeline `LokiPipeline`. -fn extract_item(content_type: ContentType, bytes: Bytes) -> Result>> +async fn extract_item( + content_type: ContentType, + bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, +) -> Result>> where LokiMiddleItem: Into, LokiMiddleItem>: Into, @@ -285,7 +298,9 @@ where LokiJsonParser::from_bytes(bytes)?.flat_map(|item| item.into_iter().map(|i| i.into())), )), x if x == *PB_CONTENT_TYPE => Ok(Box::new( - LokiPbParser::from_bytes(bytes)?.flat_map(|item| item.into_iter().map(|i| i.into())), + LokiPbParser::from_bytes(bytes, memory_limiter) + .await? + .flat_map(|item| item.into_iter().map(|i| i.into())), )), _ => UnsupportedContentTypeSnafu { content_type }.fail(), } @@ -554,8 +569,13 @@ pub struct LokiPbParser { } impl LokiPbParser { - pub fn from_bytes(bytes: Bytes) -> Result { - let decompressed = snappy_decompress_loki_request(&bytes)?; + pub async fn from_bytes( + bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, + ) -> Result { + // `decompressed` carries the memory permits for the decoded bytes and + // keeps them alive until the protobuf decoding below is finished. + let decompressed = snappy_decompress_loki_request(&bytes, memory_limiter).await?; let req = loki_proto::logproto::PushRequest::decode(&decompressed[..]) .context(DecodeLokiRequestSnafu)?; @@ -565,13 +585,34 @@ impl LokiPbParser { } } -fn snappy_decompress_loki_request(buf: &[u8]) -> Result> { +async fn snappy_decompress_loki_request( + buf: &[u8], + limiter: &crate::request_memory_limiter::ServerMemoryLimiter, +) -> Result { // Loki's protobuf push body is Snappy-compressed independent of HTTP // content-encoding, so keep this decode step explicit. + // + // The decoded size is validated before the output buffer is allocated: + // the raw Snappy header carries the decoded length, and without this cap + // a tiny body could declare multiple GiB of output. The declared size is + // also charged to the aggregate request-memory limiter before the + // allocation, so concurrent Loki pushes cannot bypass the quota. The + // returned buffer carries the permits, keeping the reservation alive + // until the caller is done with the decoded bytes. + let decoded_len = snap::raw::decompress_len(buf).context(DecompressSnappyLokiRequestSnafu)?; + ensure!( + decoded_len <= crate::prom_store::MAX_DECOMPRESSED_REQUEST_SIZE, + crate::error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: crate::prom_store::MAX_DECOMPRESSED_REQUEST_SIZE as u64, + } + ); + let guard = limiter.acquire(decoded_len as u64).await?; let mut decoder = Decoder::new(); - decoder + let data = decoder .decompress_vec(buf) - .context(DecompressSnappyLokiRequestSnafu) + .context(DecompressSnappyLokiRequestSnafu)?; + Ok(crate::prom_store::ChargedBuffer::new(data, vec![guard])) } impl Iterator for LokiPbParser { @@ -819,6 +860,7 @@ mod tests { use super::*; use crate::error::Error::{DecompressSnappyLokiRequest, InvalidLokiLabels}; use crate::prom_store::snappy_compress; + use crate::request_memory_limiter::ServerMemoryLimiter; const JSON_PAYLOAD: &[u8] = br#"{ "streams": [ @@ -869,13 +911,15 @@ mod tests { assert_eq!(prost_ts_to_nano(&ts), 1731748568804293888); } - #[test] - fn test_json_direct_ingest_builds_schema_and_pads_rows() { + #[tokio::test] + async fn test_json_direct_ingest_builds_schema_and_pads_rows() { let request = build_loki_raw_insert_request( JSON_CONTENT_TYPE.clone(), "custom_loki".to_string(), Bytes::from_static(JSON_PAYLOAD), + &ServerMemoryLimiter::default(), ) + .await .unwrap(); assert_eq!(request.table_name, "custom_loki"); @@ -916,12 +960,14 @@ mod tests { assert_eq!(row_string_value(second, 5), Some("worker-0")); } - #[test] - fn test_json_pipeline_conversion_names_loki_fields() { + #[tokio::test] + async fn test_json_pipeline_conversion_names_loki_fields() { let items = extract_item::( JSON_CONTENT_TYPE.clone(), Bytes::from_static(JSON_PAYLOAD), + &ServerMemoryLimiter::default(), ) + .await .unwrap() .collect::>(); @@ -951,8 +997,8 @@ mod tests { ); } - #[test] - fn test_protobuf_parser_decodes_snappy_push_request() { + #[tokio::test] + async fn test_protobuf_parser_decodes_snappy_push_request() { let request = PushRequest { streams: vec![StreamAdapter { labels: r#"{job="api"}"#.to_string(), @@ -973,9 +1019,14 @@ mod tests { }; let bytes = snappy_compress(&request.encode_to_vec()).unwrap(); - let items = extract_item::(PB_CONTENT_TYPE.clone(), Bytes::from(bytes)) - .unwrap() - .collect::>(); + let items = extract_item::( + PB_CONTENT_TYPE.clone(), + Bytes::from(bytes), + &ServerMemoryLimiter::default(), + ) + .await + .unwrap() + .collect::>(); assert_eq!(items.len(), 1); assert_eq!(items[0].ts, 1731748568804293888); @@ -987,9 +1038,14 @@ mod tests { assert!(!items[0].structured_metadata.is_empty()); } - #[test] - fn test_protobuf_parser_rejects_invalid_snappy_payload() { - let err = match LokiPbParser::from_bytes(Bytes::from_static(b"not-snappy")) { + #[tokio::test] + async fn test_protobuf_parser_rejects_invalid_snappy_payload() { + let err = match LokiPbParser::from_bytes( + Bytes::from_static(b"not-snappy"), + &ServerMemoryLimiter::default(), + ) + .await + { Ok(_) => panic!("expected invalid snappy payload to fail"), Err(err) => err, }; @@ -997,6 +1053,51 @@ mod tests { assert!(matches!(err, DecompressSnappyLokiRequest { .. })); } + #[tokio::test] + async fn test_protobuf_push_charges_decoded_size() { + use common_memory_manager::OnExhaustedPolicy; + + // A valid snappy payload whose decoded size exceeds the 1 KiB quota. + let request = loki_proto::logproto::PushRequest { + streams: vec![loki_proto::logproto::StreamAdapter { + labels: "{job=\"quota\"}".to_string(), + entries: std::iter::repeat_with(|| loki_proto::logproto::EntryAdapter { + timestamp: Some(Timestamp { + seconds: 1, + nanos: 0, + }), + line: "x".repeat(256), + ..Default::default() + }) + .take(8) + .collect(), + ..Default::default() + }], + }; + let bytes = Bytes::from(snappy_compress(&request.encode_to_vec()).unwrap()); + + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let err = match LokiPbParser::from_bytes(bytes.clone(), &limiter).await { + Ok(_) => panic!("expected quota exhaustion to fail"), + Err(err) => err, + }; + assert!(matches!( + err, + crate::error::Error::MemoryLimitExceeded { .. } + )); + assert_eq!(0, limiter.used_bytes(), "failed charge must be released"); + + // The same payload fits an adequately sized quota. + let limiter = ServerMemoryLimiter::new(64 * 1024, OnExhaustedPolicy::Fail); + let parser = LokiPbParser::from_bytes(bytes, &limiter).await.unwrap(); + assert_eq!(1, parser.streams.len()); + assert_eq!( + 0, + limiter.used_bytes(), + "guards must be released after decode" + ); + } + #[test] fn test_parse_loki_labels() { let mut expected = BTreeMap::new(); diff --git a/src/servers/src/http/memory_limit.rs b/src/servers/src/http/memory_limit.rs index 50b4c68f03e..41905616f44 100644 --- a/src/servers/src/http/memory_limit.rs +++ b/src/servers/src/http/memory_limit.rs @@ -13,13 +13,47 @@ // limitations under the License. //! Middleware for limiting total memory usage of concurrent HTTP request bodies. +//! +//! Admission works in two stages: +//! +//! 1. Upfront, the request's `Content-Length` (if present) is charged to the +//! shared [`ServerMemoryLimiter`], preserving the configured wait/fail +//! policy for regular requests. +//! 2. While the body is streamed, [`AccountedBody`] charges every byte beyond +//! the upfront reservation. This closes two gaps of a header-only +//! reservation: HTTP/1.1 chunked requests carry no `Content-Length` (and +//! would otherwise be admitted for free), and a client can understate the +//! header while sending a much larger body. +//! +//! Permits acquired while streaming are held by [`BodyMemoryAccounting`], +//! which is also inserted into the request extensions so they stay alive until +//! the whole request is finished (extractors collect the body into owned +//! buffers that outlive the body stream itself). +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use axum::body::Body; use axum::extract::{Request, State}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; +use bytes::Bytes; +use common_memory_manager::MemoryGuard; use http::StatusCode; +use http_body::{Body as HttpBody, Frame}; +use crate::error::Result; use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; + +/// Marker inserted by [`memory_limit_middleware`] when the request arrives +/// with a non-identity `Content-Encoding` header. Route-local +/// [`decoded_body_accounting_middleware`] uses it to account the decompressed +/// body: only such requests expand during request decompression. +#[derive(Clone, Copy)] +pub(crate) struct ContentEncoded; pub async fn memory_limit_middleware( State(limiter): State, @@ -44,5 +78,571 @@ pub async fn memory_limit_middleware( } }; - next.run(req).await + let content_encoded = req + .headers() + .get(http::header::CONTENT_ENCODING) + .and_then(|v| v.to_str().ok()) + .is_some_and(|v| !v.eq_ignore_ascii_case("identity")); + + // Account whatever is actually streamed beyond the upfront reservation + // (chunked transfer has none). + let accounting = BodyMemoryAccounting::default(); + // The middleware keeps its own handle so the acquired permits stay held + // for the whole request: the body wrapper and the extensions are dropped + // once the extractors finish collecting the body, before the handler is + // done with the decoded data. + let retained_accounting = accounting.clone(); + let (mut parts, body) = req.into_parts(); + // Expose the limiter to handlers that decompress request bodies on their + // own (e.g. the Loki protobuf push) so they can charge the decoded size. + parts.extensions.insert(limiter.clone()); + parts.extensions.insert(accounting.clone()); + if content_encoded { + parts.extensions.insert(ContentEncoded); + } + let accounted = AccountedBody::new(body, limiter, content_length, accounting); + let req = Request::from_parts(parts, Body::new(accounted)); + + let response = next.run(req).await; + // `retained_accounting` stays alive until this function returns, holding + // the permits across the handler's use of the collected body. + quota_exceeded_response(&retained_accounting, response) +} + +/// Rewrites the response of a request whose body accounting hit the quota: +/// the failure aborts the body mid-stream, so the extractors reject the +/// request with a generic body error (mapped to 400). Quota exhaustion must +/// surface as 429, matching the upfront admission path, so clients can tell +/// backpressure apart from malformed input. +fn quota_exceeded_response(accounting: &BodyMemoryAccounting, response: Response) -> Response { + // The flag can only be set when the body failed before the handler ran, + // so the response is the extractor rejection and can be replaced. + if accounting.take_quota_exceeded() { + return ( + StatusCode::TOO_MANY_REQUESTS, + "Request body memory limit exceeded", + ) + .into_response(); + } + response +} + +/// Route-local accounting for routes that install `RequestDecompressionLayer`: +/// it must be layered *inside* the +/// decompression layer, so the body it wraps is the **decompressed** stream. +/// It requires the global accounting installed by [`crate::http::HttpServer`] +/// and must use the same [`ServerMemoryLimiter`]. +/// +/// Only encoded requests marked by the global middleware are +/// accounted; plain (uncompressed) requests are already covered by the global +/// wire-byte accounting, and charging them here as well would double-count. +/// For decompressed requests the `Content-Length` header has been stripped by +/// the decompression layer, so the full decoded size is charged as it streams. +pub async fn decoded_body_accounting_middleware( + State(limiter): State, + req: Request, + next: Next, +) -> Response { + if req.extensions().get::().is_none() { + return next.run(req).await; + } + + let content_length = req + .headers() + .get(http::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + .unwrap_or(0); + + let accounting = BodyMemoryAccounting::default(); + // Retain the permits for the whole request, like the global middleware. + let retained_accounting = accounting.clone(); + let (mut parts, body) = req.into_parts(); + parts.extensions.insert(accounting.clone()); + let accounted = AccountedBody::new(body, limiter, content_length, accounting); + let req = Request::from_parts(parts, Body::new(accounted)); + + let response = next.run(req).await; + // `retained_accounting` stays alive until this function returns, holding + // the permits across the handler's use of the collected body. + quota_exceeded_response(&retained_accounting, response) +} + +/// Holds the memory guards acquired while a request body is streamed. Shared +/// between the [`AccountedBody`] wrapper, the request extensions and the +/// middleware itself, so the permits are only released when the request +/// (including any collected body buffers and the handler still using them) +/// is finished. +#[derive(Clone, Default)] +struct BodyMemoryAccounting { + guards: Arc>>>, + /// Set when an incremental charge hit the quota, so the rejection can be + /// rewritten into a 429 after the middleware regains control. + quota_exceeded: Arc, +} + +impl BodyMemoryAccounting { + fn hold(&self, guard: MemoryGuard) { + self.guards.lock().unwrap().push(guard); + } + + fn mark_quota_exceeded(&self) { + self.quota_exceeded.store(true, Ordering::Release); + } + + fn take_quota_exceeded(&self) -> bool { + self.quota_exceeded.swap(false, Ordering::AcqRel) + } +} + +type AcquireFuture = + Pin>> + Send>>; + +enum ChargeOutcome { + Charged, + Pending, + Failed, +} + +/// A request body wrapper that charges the shared limiter for every byte +/// streamed beyond `pre_charged` (the upfront `Content-Length` reservation). +struct AccountedBody { + inner: Body, + limiter: ServerMemoryLimiter, + accounting: BodyMemoryAccounting, + /// Bytes of body data streamed so far. + streamed: u64, + /// Bytes of the streamed body already covered by incremental permits. + incrementally_charged: u64, + /// Bytes covered by the upfront `Content-Length` reservation. + pre_charged: u64, + /// A data frame held back while its permit is being acquired, together + /// with the number of extra bytes to charge for it. + pending: Option<(Frame, u64)>, + pending_acquire: Option, +} + +impl AccountedBody { + fn new( + inner: Body, + limiter: ServerMemoryLimiter, + pre_charged: u64, + accounting: BodyMemoryAccounting, + ) -> Self { + Self { + inner, + limiter, + accounting, + streamed: 0, + incrementally_charged: 0, + pre_charged, + pending: None, + pending_acquire: None, + } + } + + /// Bytes that still need an incremental permit for the body streamed so far. + fn uncharged(&self) -> u64 { + self.streamed + .saturating_sub(self.pre_charged) + .saturating_sub(self.incrementally_charged) + } + + /// Acquires a permit for `bytes`; only one acquisition may be in flight. + /// + /// A parked acquisition always belongs to the currently buffered frame, so + /// it must be resolved *before* any new acquisition is tried: letting the + /// fast path succeed while a waiter is parked would leave the waiter alive, + /// and its late completion would then be credited with another frame's + /// byte count (under-reserving memory relative to what is marked charged). + fn charge(&mut self, bytes: u64, cx: &mut Context<'_>) -> ChargeOutcome { + debug_assert!(bytes > 0); + if let Some(fut) = self.pending_acquire.as_mut() { + return match fut.as_mut().poll(cx) { + Poll::Ready(Ok(guard)) => { + self.accounting.hold(guard); + self.pending_acquire = None; + self.incrementally_charged += bytes; + ChargeOutcome::Charged + } + Poll::Ready(Err(_)) => { + self.pending_acquire = None; + ChargeOutcome::Failed + } + Poll::Pending => ChargeOutcome::Pending, + }; + } + // Fast path: the permit is immediately available. + if let Some(guard) = self.limiter.try_acquire(bytes) { + self.accounting.hold(guard); + self.incrementally_charged += bytes; + return ChargeOutcome::Charged; + } + let limiter = self.limiter.clone(); + let mut fut = Box::pin(async move { limiter.acquire(bytes).await }); + // Poll immediately so a ready result is not missed (returning + // `Pending` without registering progress would park the task forever). + match fut.as_mut().poll(cx) { + Poll::Ready(Ok(guard)) => { + self.accounting.hold(guard); + self.incrementally_charged += bytes; + ChargeOutcome::Charged + } + Poll::Ready(Err(_)) => ChargeOutcome::Failed, + Poll::Pending => { + self.pending_acquire = Some(fut); + ChargeOutcome::Pending + } + } + } + + /// Hands out `frame` once `bytes` of it are charged, buffering the frame + /// while the acquisition is pending. + fn hand_out( + &mut self, + frame: Frame, + bytes: u64, + cx: &mut Context<'_>, + ) -> Poll, axum::Error>>> { + if bytes == 0 { + return Poll::Ready(Some(Ok(frame))); + } + match self.charge(bytes, cx) { + ChargeOutcome::Charged => Poll::Ready(Some(Ok(frame))), + ChargeOutcome::Pending => { + self.pending = Some((frame, bytes)); + Poll::Pending + } + ChargeOutcome::Failed => { + self.accounting.mark_quota_exceeded(); + Poll::Ready(Some(Err(limit_exceeded_error()))) + } + } + } +} + +impl HttpBody for AccountedBody { + type Data = Bytes; + type Error = axum::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + + // Finish a pending acquisition for a buffered frame first. + if let Some((frame, bytes)) = this.pending.take() { + return this.hand_out(frame, bytes, cx); + } + + match Pin::new(&mut this.inner).poll_frame(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => Poll::Ready(None), + Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), + Poll::Ready(Some(Ok(frame))) => { + if let Some(data) = frame.data_ref() { + this.streamed += data.len() as u64; + let uncharged = this.uncharged(); + this.hand_out(frame, uncharged, cx) + } else { + // Trailers pass through uncharged. + Poll::Ready(Some(Ok(frame))) + } + } + } + } + + fn is_end_stream(&self) -> bool { + self.pending.is_none() && self.inner.is_end_stream() + } + + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +fn limit_exceeded_error() -> axum::Error { + axum::Error::new(std::io::Error::other( + "request body exceeded the aggregate memory limit", + )) +} + +#[cfg(test)] +mod tests { + use axum::body::{Body, Bytes}; + use axum::http::{Request, StatusCode, header}; + use axum::routing::post; + use axum::{Router, middleware}; + use common_memory_manager::OnExhaustedPolicy; + use tower::ServiceExt; + + use super::memory_limit_middleware; + use crate::request_memory_limiter::ServerMemoryLimiter; + + /// Streams `body` in `chunk_size` frames without a Content-Length header + /// (like HTTP/1.1 chunked transfer). + fn chunked_request(uri: &str, body: Vec, chunk_size: usize) -> Request { + let stream = futures_util::stream::iter( + body.chunks(chunk_size) + .map(|c| Ok::<_, std::io::Error>(Bytes::copy_from_slice(c))) + .collect::>(), + ); + Request::builder() + .uri(uri) + .method("POST") + .body(Body::from_stream(stream)) + .unwrap() + } + + fn counted_body_app(limiter: ServerMemoryLimiter) -> Router { + Router::new() + .route( + "/echo", + post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }), + ) + .layer(middleware::from_fn_with_state( + limiter, + memory_limit_middleware, + )) + } + + #[tokio::test] + async fn test_chunked_body_larger_than_quota_is_rejected() { + // The quota cannot cover a fully streamed 64 KiB chunked body. + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let res = app + .oneshot(chunked_request("/echo", vec![b'x'; 64 * 1024], 1024)) + .await + .unwrap(); + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "a chunked body larger than the quota must be rejected with 429, not 400" + ); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_small_chunked_body_fits() { + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let res = app + .oneshot(chunked_request("/echo", vec![b'x'; 1024], 512)) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], b"1024"); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_content_length_still_admitted_upfront() { + // A Content-Length larger than the quota is rejected before the body + // is read (429), unchanged from the previous behavior. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "4096") + .body(Body::from(vec![b'x'; 4096])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn test_understated_content_length_is_charged_the_difference() { + // Claim 1 KiB in Content-Length but stream 16 KiB: the extra 15 KiB + // must be charged while streaming and rejected by the small quota. + let limiter = ServerMemoryLimiter::new(2 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let stream = futures_util::stream::iter(vec![Ok::<_, std::io::Error>( + Bytes::copy_from_slice(&vec![b'x'; 16 * 1024]), + )]); + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "1024") + .body(Body::from_stream(stream)) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_honest_content_length_is_not_double_charged() { + // With a quota large enough for the whole body, an accurate + // Content-Length must admit the request without extra charges. + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "4096") + .body(Body::from(vec![b'x'; 4096])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(0, limiter.used_bytes()); + } + + /// Mirrors the route-local stack: decoded-body accounting *inside* a + /// request-decompression layer, like the ingestion routes in `http.rs`. + fn decompressing_app(limiter: ServerMemoryLimiter) -> Router { + use tower_http::decompression::RequestDecompressionLayer; + + Router::new() + .route( + "/echo", + post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }), + ) + .layer(middleware::from_fn_with_state( + limiter, + super::decoded_body_accounting_middleware, + )) + .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)) + } + + fn zstd_body(decoded: &[u8]) -> Body { + let compressed = zstd::stream::encode_all(decoded, 3).unwrap(); + Body::from(compressed) + } + + #[tokio::test] + async fn test_decompressed_body_is_charged_for_content_encoded_requests() { + // The decoded size (64 KiB) exceeds the 16 KiB quota: the accounting + // layer must see the decompressed bytes and reject mid-stream. + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .extension(super::ContentEncoded) + .body(zstd_body(&vec![b'x'; 64 * 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the quota must be rejected with 429" + ); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_small_decompressed_body_fits() { + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .extension(super::ContentEncoded) + .body(zstd_body(&vec![b'x'; 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], b"1024", "the handler must see the decoded body"); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_pending_acquisition_is_attributed_to_its_own_bytes() { + use std::time::Duration; + + use tokio_stream::wrappers::ReceiverStream; + + // Regression test: a parked acquisition must be resolved before any + // new acquisition is tried; otherwise a late waiter completion can be + // credited with a later frame's byte count. + // + // Scenario: quota 8 KiB fully held externally; frame A (4 KiB) parks + // its acquisition; the external guard is then released so A drains; + // frame B (8 KiB) arrives. Correct behavior: B waits for its own 8 KiB + // acquisition and is rejected when it times out (only 4 KiB are free). + // Buggy behavior: the stale waiter (4 KiB) completes and is credited + // with B's 8 KiB, so B is delivered although the quota cannot cover it. + let limiter = ServerMemoryLimiter::new( + 8 * 1024, + OnExhaustedPolicy::Wait { + timeout: Duration::from_millis(500), + }, + ); + let app = counted_body_app(limiter.clone()); + + let (tx, rx) = tokio::sync::mpsc::channel::>(2); + let driver_limiter = limiter.clone(); + let driver = tokio::spawn(async move { + // Hold the whole quota so frame A's acquisition parks. + let external = driver_limiter.acquire(8 * 1024).await.unwrap(); + tx.send(Ok(Bytes::from(vec![b'a'; 4096]))).await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + // Free the quota: frame A's waiter drains and A is delivered. + drop(external); + tokio::time::sleep(Duration::from_millis(300)).await; + tx.send(Ok(Bytes::from(vec![b'b'; 8192]))).await.unwrap(); + // Keep the channel open; frame B's own acquisition must time out. + tokio::time::sleep(Duration::from_millis(1500)).await; + }); + + let res = app + .oneshot( + Request::builder() + .uri("/echo") + .method("POST") + .body(Body::from_stream(ReceiverStream::new(rx))) + .unwrap(), + ) + .await + .unwrap(); + driver.await.unwrap(); + + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "frame B must wait for its own 8 KiB and time out, not borrow A's reservation" + ); + } + + #[tokio::test] + async fn test_plain_requests_are_not_charged_by_the_decoded_layer() { + // Without the ContentEncoded marker the decoded layer must be a + // no-op: plain bodies are accounted by the global wire accounting. + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + // Plain body, no marker: passes through. + let req = Request::builder() + .uri("/echo") + .method("POST") + .body(Body::from(vec![b'x'; 1024])) + .unwrap(); + let res = app.clone().oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + // Compressed body WITHOUT the marker (e.g. a route that has the + // decompression layer but no accounting): still passes through. + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .body(zstd_body(&vec![b'x'; 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + } } diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index 2d5385c5e33..717e6ca65da 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -77,7 +77,6 @@ fn content_type_to_string(content_type: Option<&TypedHeader>) -> St #[derive(Clone)] pub struct OtlpState { pub with_metric_engine: bool, - pub experimental_enable_exponential_histogram: bool, pub handler: OpenTelemetryProtocolHandlerRef, } @@ -108,7 +107,6 @@ pub async fn metrics( let OtlpState { with_metric_engine, - experimental_enable_exponential_histogram, handler, } = state; @@ -117,7 +115,6 @@ pub async fn metrics( resource_attrs: http_opts.resource_attrs, promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, - experimental_enable_exponential_histogram, // set by the frontend from its config is_legacy: false, resource_info: false, diff --git a/src/servers/src/http/pprof/nix.rs b/src/servers/src/http/pprof/nix.rs deleted file mode 100644 index 12a518e1de0..00000000000 --- a/src/servers/src/http/pprof/nix.rs +++ /dev/null @@ -1,123 +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::time::Duration; - -use common_error::ext::ErrorExt; -use common_error::status_code::StatusCode; -use common_macro::stack_trace_debug; -use prost::Message; -use snafu::{Location, ResultExt, Snafu}; - -#[derive(Snafu)] -#[stack_trace_debug] -pub enum Error { - #[snafu(display("Failed to create profiler guard"))] - CreateGuard { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to create report"))] - CreateReport { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to create flamegraph"))] - CreateFlamegraph { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to create pprof report"))] - ReportPprof { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, -} - -pub type Result = std::result::Result; - -impl ErrorExt for Error { - fn status_code(&self) -> StatusCode { - StatusCode::Unexpected - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -/// CPU profiler utility. -// Inspired by https://github.com/datafuselabs/databend/blob/67f445e83cd4eceda98f6c1c114858929d564029/src/common/base/src/base/profiling.rs -#[derive(Debug)] -pub struct Profiling { - /// Sample duration. - duration: Duration, - /// Sample frequency. - frequency: i32, -} - -impl Profiling { - /// Creates a new profiler. - pub fn new(duration: Duration, frequency: i32) -> Profiling { - Profiling { - duration, - frequency, - } - } - - /// Profiles and returns a generated pprof report. - pub async fn report(&self) -> Result { - let guard = pprof::ProfilerGuardBuilder::default() - .frequency(self.frequency) - .blocklist(&["libc", "libgcc", "pthread", "vdso"]) - .build() - .context(CreateGuardSnafu)?; - tokio::time::sleep(self.duration).await; - guard.report().build().context(CreateReportSnafu) - } - - /// Profiles and returns a generated flamegraph. - pub async fn dump_flamegraph(&self) -> Result> { - let mut body: Vec = Vec::new(); - - let report = self.report().await?; - report - .flamegraph(&mut body) - .context(CreateFlamegraphSnafu)?; - - Ok(body) - } - - /// Profiles and returns a generated proto. - pub async fn dump_proto(&self) -> Result> { - let report = self.report().await?; - // Generate google’s pprof format report. - let profile = report.pprof().context(ReportPprofSnafu)?; - let body = profile.encode_to_vec(); - - Ok(body) - } -} diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 17db6abae04..05d9725f947 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -52,8 +52,9 @@ use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_remote_write::{ REMOTE_WRITE_V1_VERSION, REMOTE_WRITE_V2_VERSION, decode_remote_write_request, }; -use crate::prom_store::snappy_decompress; +use crate::prom_store::{MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited}; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; +use crate::request_memory_limiter::ServerMemoryLimiter; pub const PHYSICAL_TABLE_PARAM: &str = "physical_table"; pub const DEFAULT_ENCODING: &str = "snappy"; @@ -74,8 +75,10 @@ pub struct PromStoreState { pub pipeline_handler: Option, pub prom_store_with_metric_engine: bool, pub prom_validation_mode: PromValidationMode, - pub experimental_enable_prometheus_native_histogram: bool, pub pending_rows_batcher: Option>, + /// Shared request-memory limiter used to charge decompressed remote + /// read/write bodies against the aggregate quota. + pub memory_limiter: ServerMemoryLimiter, } #[derive(Debug, Serialize, Deserialize)] @@ -146,8 +149,8 @@ async fn remote_write_v1( pipeline_handler, prom_store_with_metric_engine, prom_validation_mode, - experimental_enable_prometheus_native_histogram: _, pending_rows_batcher, + memory_limiter, } = state; if let Some(response) = vm_proto_version_response(¶ms) { @@ -175,13 +178,22 @@ async fn remote_write_v1( processor.set_pipeline(pipeline_handler, query_ctx.clone(), pipeline_def); } - let mut decoded = - decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?; + // Keep the decode guards alive until the batches are built: the + // `TablesBuilder` retains the decompressed buffer as its raw data. + let (mut decoded, decode_guards) = decode_remote_write_request( + is_zstd, + body, + prom_validation_mode, + &mut processor, + &memory_limiter, + ) + .await?; // Parsing borrows the decode buffer, but row building copies out of it: tag // values through `decode_string`, column names through `to_owned`, and the // borrowing `col_indexes` dies inside `as_insert_requests`. Nothing below - // references the buffer, so it need not span the write. + // references the buffer, so it need not span the write — the same goes for + // the decode guards charged against the aggregate memory quota. let req = if processor.use_pipeline { drop(decoded); processor.exec_pipeline().await? @@ -190,6 +202,7 @@ async fn remote_write_v1( drop(decoded); req }; + drop(decode_guards); let batches = into_prom_write_batches(req, query_ctx); let outcome = match write_prometheus_rows_with_progress( @@ -228,8 +241,8 @@ async fn remote_write_v2( pipeline_handler: _, prom_store_with_metric_engine, prom_validation_mode: _, - experimental_enable_prometheus_native_histogram, pending_rows_batcher, + memory_limiter, } = state; if let Some(response) = vm_proto_version_response(¶ms) { @@ -243,11 +256,7 @@ async fn remote_write_v2( let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); - let req = match decode_remote_write_v2( - is_zstd, - body, - experimental_enable_prometheus_native_histogram, - ) { + let req = match decode_remote_write_v2(is_zstd, body, &memory_limiter).await { Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; @@ -404,27 +413,45 @@ async fn write_prometheus_rows_with_progress( mut batches: Vec, ) -> std::result::Result { if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher { + // Preflight before the bulk eligibility decision: pre_write hooks + // may redirect the contexts (e.g. to a per-tenant schema), so + // eligibility must be evaluated against the prepared batches — + // and the fallback below reuses them without re-running hooks or + // admission. preflight_prometheus_rows(&prom_store_handler, &mut batches) .await .map_err(|error| PromWriteError { error, rows_written: 0, })?; - let mut rows_written = 0; - for (temp_ctx, reqs) in batches { - let rows = batcher - .submit(reqs, temp_ctx) - .await - .map_err(|error| PromWriteError { - error, - rows_written, - })?; - rows_written += rows; + // Destinations bound to another physical table must stay on the + // ordinary insert path, which routes per destination; time index + // units need no check — the bulk encode converts each request to + // its destination's unit. + if batcher.accepts_bulk_destinations(batches.iter()).await { + let mut rows_written = 0; + for (temp_ctx, reqs) in batches { + let rows = + batcher + .submit(reqs, temp_ctx) + .await + .map_err(|error| PromWriteError { + error, + rows_written, + })?; + rows_written += rows; + } + return Ok(PromWriteOutcome { + write_cost: 0, + rows_written, + }); } - return Ok(PromWriteOutcome { - write_cost: 0, - rows_written, - }); + return write_prepared_prometheus_rows_with_progress( + prom_store_handler, + batches, + prom_store_with_metric_engine, + ) + .await; } let row_counts = batches @@ -463,6 +490,33 @@ async fn write_prometheus_rows_with_progress( }) } +/// Writes already-preflighted batches through the ordinary (prepared) write +/// path; the pre_write hooks and admission have already run. +async fn write_prepared_prometheus_rows_with_progress( + prom_store_handler: PromStoreProtocolHandlerRef, + batches: Vec, + prom_store_with_metric_engine: bool, +) -> std::result::Result { + let mut write_cost = 0; + let mut rows_written = 0; + for (ctx, request) in batches { + let rows = prom_write_row_count(&request); + let output = prom_store_handler + .write_prepared(request, ctx, prom_store_with_metric_engine) + .await + .map_err(|error| PromWriteError { + error, + rows_written, + })?; + write_cost += output.meta.cost; + rows_written += rows; + } + Ok(PromWriteOutcome { + write_cost, + rows_written, + }) +} + async fn write_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, pending_rows_batcher: Option>, @@ -490,20 +544,44 @@ async fn write_prometheus_v2_rows_with_progress( }); } + let sample_batch_count = sample_batches.len(); + let mut batches = sample_batches; + batches.extend(histogram_batches); + if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher { - return write_batched_prometheus_v2_rows_with_progress( + // Same ordering as the v1 path: preflight (which may redirect the + // contexts) before the bulk eligibility decision, and the fallback + // reuses the prepared batches without re-running hooks or admission. + preflight_prometheus_rows(&prom_store_handler, &mut batches) + .await + .map_err(|error| PromWriteV2Error { + error, + samples_written: 0, + histograms_written: 0, + })?; + // Destinations bound to another physical table must stay on the + // ordinary insert path, which routes per destination; time index + // units need no check — the bulk encode converts each request to + // its destination's unit. + if batcher.accepts_bulk_destinations(batches.iter()).await { + return write_batched_prometheus_v2_rows_with_progress( + prom_store_handler, + batcher.as_ref(), + prom_store_with_metric_engine, + sample_batch_count, + batches, + ) + .await; + } + return write_prepared_prometheus_v2_rows_with_progress( prom_store_handler, - batcher.as_ref(), + batches, + sample_batch_count, prom_store_with_metric_engine, - sample_batches, - histogram_batches, ) .await; } - let sample_batch_count = sample_batches.len(); - let mut batches = sample_batches; - batches.extend(histogram_batches); let row_counts = batches .iter() .map(|(_, request)| prom_write_row_count(request)) @@ -551,24 +629,49 @@ async fn write_prometheus_v2_rows_with_progress( }) } +/// Writes already-preflighted batches through the ordinary (prepared) write +/// path with v2 partial-progress accounting; the pre_write hooks and +/// admission have already run. +async fn write_prepared_prometheus_v2_rows_with_progress( + prom_store_handler: PromStoreProtocolHandlerRef, + batches: Vec, + sample_batch_count: usize, + prom_store_with_metric_engine: bool, +) -> std::result::Result { + let mut write_cost = 0; + let mut samples_written = 0; + let mut histograms_written = 0; + for (index, (ctx, request)) in batches.into_iter().enumerate() { + let rows = prom_write_row_count(&request); + let output = prom_store_handler + .write_prepared(request, ctx, prom_store_with_metric_engine) + .await + .map_err(|error| PromWriteV2Error { + error, + samples_written, + histograms_written, + })?; + write_cost += output.meta.cost; + if index < sample_batch_count { + samples_written += rows; + } else { + histograms_written += rows; + } + } + Ok(PromWriteV2Outcome { + write_cost, + samples_written, + histograms_written, + }) +} + async fn write_batched_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, batcher: &B, prom_store_with_metric_engine: bool, - sample_batches: Vec, - histogram_batches: Vec, + sample_batch_count: usize, + batches: Vec, ) -> std::result::Result { - let sample_batch_count = sample_batches.len(); - let mut batches = sample_batches; - batches.extend(histogram_batches); - preflight_prometheus_rows(&prom_store_handler, &mut batches) - .await - .map_err(|error| PromWriteV2Error { - error, - samples_written: 0, - histograms_written: 0, - })?; - let mut samples_written = 0; let mut histograms_written = 0; let mut write_cost = 0; @@ -753,7 +856,7 @@ pub async fn remote_read( let db = params.db.clone().unwrap_or_default(); query_ctx.set_channel(Channel::Prometheus); - let request = decode_remote_read_request(body).await?; + let request = decode_remote_read_request(body, &state.memory_limiter).await?; let query_ctx = Arc::new(query_ctx); let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_READ_ELAPSED @@ -763,8 +866,13 @@ pub async fn remote_read( state.prom_store_handler.read(request, query_ctx).await } -async fn decode_remote_read_request(body: Bytes) -> Result { - let buf = snappy_decompress(&body[..])?; +async fn decode_remote_read_request( + body: Bytes, + limiter: &ServerMemoryLimiter, +) -> Result { + // Holds the memory permits for the decompressed bytes until the protobuf + // decoding is finished. + let buf = snappy_decompress_limited(&body[..], MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await?; ReadRequest::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } @@ -869,14 +977,18 @@ mod tests { events: events.clone(), }; - let Ok(outcome) = write_batched_prometheus_v2_rows_with_progress( - handler, - &batcher, - true, - vec![test_prom_write_batch("sample")], - vec![test_prom_write_batch("histogram")], - ) - .await + let mut batches = vec![ + test_prom_write_batch("sample"), + test_prom_write_batch("histogram"), + ]; + // The caller preflights before choosing the bulk path; the batched + // writer consumes the prepared batches. + preflight_prometheus_rows(&handler, &mut batches) + .await + .unwrap(); + let Ok(outcome) = + write_batched_prometheus_v2_rows_with_progress(handler, &batcher, true, 1, batches) + .await else { panic!("mixed remote write should succeed") }; @@ -1017,8 +1129,8 @@ mod tests { pipeline_handler: None, prom_store_with_metric_engine: false, prom_validation_mode: PromValidationMode::Strict, - experimental_enable_prometheus_native_histogram: false, pending_rows_batcher: None, + memory_limiter: ServerMemoryLimiter::default(), } } diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 420a0c8f4b4..e417044ef43 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -49,11 +49,11 @@ use futures::future::join_all; use futures::{StreamExt, TryStreamExt}; use itertools::Itertools; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; -use promql_parser::parser::token::{self}; +use promql_parser::parser::token::TokenType; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ - AggregateExpr, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, MatrixSelector, ParenExpr, - SubqueryExpr, UnaryExpr, VectorSelector, + AggregateExpr, BinModifier, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, + MatrixSelector, ParenExpr, SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector, }; use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement}; use serde::de::{self, MapAccess, Visitor}; @@ -547,6 +547,12 @@ pub async fn range_query( StatusCode::InvalidArguments ); let promql_expr = prom_query.expr(); + // Prometheus rejects a range-vector expression in a range query even when start == end, + // which the planner cannot tell apart from an instant query. + if promql_expr.value_type() == ValueType::Matrix { + let err = query::promql::error::RangeVectorInRangeQuerySnafu.build(); + return PrometheusJsonResponse::error(err.status_code(), err.output_msg()); + } let metric_name_discovery = try_call_return_response!(find_metric_name_not_equal_matchers(promql_expr)); @@ -1373,6 +1379,72 @@ fn promql_expr_to_metric_name(expr: &PromqlExpr) -> Option { } } +/// Whether the metric name survives a binary operation; see +/// [`binary_metric_name_side`] for the operand that supplies it. +/// +/// Follows Prometheus' `shouldDropMetricName` and `resultMetric`: set operators +/// and comparisons keep the name, while arithmetic operators drop it. A +/// comparison also drops it when it returns a bool or is a one-to-one +/// `on(...)` match that does not keep `__name__` or an `ignoring(...)` match +/// that excludes it. Prometheus keeps only the `on(...)` labels or deletes the +/// `ignoring(...)` labels for one-to-one matches. +fn binary_keeps_metric_name(op: &TokenType, modifier: Option<&BinModifier>) -> bool { + if op.is_set_operator() { + return true; + } + + if !op.is_comparison_operator() { + return false; + } + + let Some(m) = modifier else { + return true; + }; + + if m.return_bool { + return false; + } + + if m.card == VectorMatchCardinality::OneToOne { + match &m.matching { + Some(LabelModifier::Include(labels)) + if !labels.labels.contains(&METRIC_NAME.to_string()) => + { + return false; + } + Some(LabelModifier::Exclude(labels)) + if labels.labels.contains(&METRIC_NAME.to_string()) => + { + return false; + } + _ => {} + } + } + + true +} + +/// The operand whose labels supply the result metric name of a binary +/// operation. +/// +/// Prometheus builds the result from the side whose labels it carries, i.e. the +/// "many" side of the match: `group_right` makes the right-hand operand that +/// side, and a scalar left-hand operand (for example `0.5 < a`) leaves the +/// vector on the right. Every other case takes the left-hand operand. +fn binary_metric_name_side<'a>( + lhs: &'a PromqlExpr, + rhs: &'a PromqlExpr, + modifier: Option<&BinModifier>, +) -> &'a PromqlExpr { + let group_right = + modifier.is_some_and(|m| matches!(m.card, VectorMatchCardinality::OneToMany(_))); + if group_right || matches!(lhs.value_type(), ValueType::Scalar) { + rhs + } else { + lhs + } +} + /// Recursively collect all metric names from a PromQL expression fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet) { match expr { @@ -1395,14 +1467,16 @@ fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet) { collect_metric_names(expr, metric_names) } PromqlExpr::Unary(UnaryExpr { .. }) => metric_names.clear(), - PromqlExpr::Binary(BinaryExpr { lhs, op, .. }) => { - if matches!( - op.id(), - token::T_LAND // INTERSECT - | token::T_LOR // UNION - | token::T_LUNLESS // EXCEPT - ) { - collect_metric_names(lhs, metric_names) + PromqlExpr::Binary(BinaryExpr { + lhs, + op, + rhs, + modifier, + .. + }) => { + if binary_keeps_metric_name(op, modifier.as_ref()) { + let side = binary_metric_name_side(lhs, rhs, modifier.as_ref()); + collect_metric_names(side, metric_names) } else { metric_names.clear() } @@ -2762,6 +2836,45 @@ mod tests { ); } + #[tokio::test] + async fn range_query_rejects_range_vector_with_equal_endpoints() { + let handler = Arc::new(TestPrometheusHandler { + catalog_manager: MemoryCatalogManager::new(), + deny_operation: false, + denied_table: None, + metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }); + let response = range_query( + State(handler.clone() as PrometheusHandlerRef), + Query(RangeQuery { + query: Some("vector(1)[5m:]".to_string()), + start: Some("300".to_string()), + end: Some("300".to_string()), + step: Some("1".to_string()), + ..Default::default() + }), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Form(RangeQuery::default()), + ) + .await; + + assert_eq!(Some(StatusCode::InvalidArguments), response.status_code); + assert_eq!( + response.error.as_deref(), + Some( + "invalid expression type \"range vector\" for range query, must be Scalar or instant Vector" + ) + ); + assert!(handler.queries.lock().unwrap().is_empty()); + } + #[tokio::test] async fn test_promql_timer_records_parse_errors() { let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { @@ -3211,6 +3324,134 @@ mod tests { expected_type: ValueType::Vector, should_error: false, }, + // Arithmetic operators always drop the metric name. + TestCase { + name: "arithmetic between metrics", + promql: "a + b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "arithmetic with scalar", + promql: "a * 2", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + // Comparisons keep the metric name from the side whose labels the + // result carries, i.e. the "many" side of the match: normally the + // left operand, but `group_right` takes it from the right operand and + // a scalar left operand leaves the vector on the right. The name is + // dropped when the comparison returns a bool or is a one-to-one + // `on(...)` match that does not keep `__name__` or an `ignoring(...)` + // match that excludes it. A `group_left`/`group_right` modifier does + // not reduce the labels, so the name is kept. + TestCase { + name: "bool comparison between metrics", + promql: "a > bool b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "bool comparison with scalar", + promql: "a > bool 0.5", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with scalar", + promql: "a > 0.5", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with scalar on the left", + promql: "0.5 < a", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "equality comparison with scalar", + promql: "up == 0", + expected_metric: Some("up"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison between metrics", + promql: "a > b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring", + promql: "a > ignoring(x) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__", + promql: "a > ignoring(__name__) b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__ and group_left", + promql: "a > ignoring(__name__) group_left b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__ and group_right", + promql: "a > ignoring(__name__) group_right b", + expected_metric: Some("b"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "unless with ignoring __name__", + promql: "a unless ignoring(__name__) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on", + promql: "a > on(x) b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on keeping __name__", + promql: "a > on(__name__) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on and group_left", + promql: "a > on(x) group_left b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on and group_right", + promql: "a > on(x) group_right b", + expected_metric: Some("b"), + expected_type: ValueType::Vector, + should_error: false, + }, // Subquery cases TestCase { name: "basic subquery", diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 518d0f65e32..05956716d4b 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -231,7 +231,7 @@ impl PrometheusJsonResponse { ) -> Self { // Hold the collector while a streaming result is consumed. let collector = query_id.and_then(get_promql_annotation_collector); - let response: Result = try { + let response: Result = async { let result = result?; let mut resp = match result.data { @@ -259,8 +259,9 @@ impl PrometheusJsonResponse { resp.resp_metrics = re; } - resp - }; + Ok(resp) + } + .await; let result_type_string = result_type.to_string(); diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 1106f89d3e9..2fed3d43de0 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -12,9 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(try_blocks)] -#![feature(exclusive_wrapper)] - use datafusion_expr::LogicalPlan; use sql::statements::statement::Statement; // Re-export for use in add_service! macro diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index 0d70540be31..888bdc42cdf 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -130,6 +130,13 @@ lazy_static! { &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] ) .unwrap(); + /// OTLP exponential-histogram data points rejected during conversion, by reason. + pub static ref OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS: IntCounterVec = register_int_counter_vec!( + "greptime_servers_otlp_exponential_histogram_rejected_data_points_total", + "Number of OTLP exponential histogram data points rejected during conversion", + &["reason"] + ) + .unwrap(); pub static ref PENDING_BATCHES: IntGauge = register_int_gauge!( "greptime_prom_store_pending_batches", "Number of pending batches waiting to be flushed" diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index 3764591beaa..bdfc7a72057 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -23,6 +23,7 @@ use async_trait::async_trait; use chrono::{NaiveDate, NaiveDateTime}; use common_catalog::parse_optional_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; +use common_frontend::metrics; use common_query::Output; use common_telemetry::{debug, error, tracing, warn}; use common_time::Timezone; @@ -85,6 +86,7 @@ pub struct MysqlInstanceShim { prepared_stmts_counter: AtomicU32, process_id: u32, prepared_stmt_cache_size: usize, + batching_enabled: bool, } impl MysqlInstanceShim { @@ -122,9 +124,22 @@ impl MysqlInstanceShim { prepared_stmts_counter: AtomicU32::new(1), process_id, prepared_stmt_cache_size, + batching_enabled: false, } } + /// Enables ordinary-table batching for this connection. + pub fn with_batching_enabled(mut self, enabled: bool) -> Self { + self.batching_enabled = enabled; + self + } + + fn new_query_context(&self) -> QueryContextRef { + let mut ctx = self.session.new_query_context(); + Arc::make_mut(&mut ctx).set_batching_enabled(self.batching_enabled); + ctx + } + #[tracing::instrument(skip_all, name = "mysql::do_query")] async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec> { if let Some(output) = @@ -168,8 +183,9 @@ impl MysqlInstanceShim { /// Retrieve the query and logical plan by a given statement key fn plan(&self, stmt_key: &str) -> Option { - let guard = self.prepared_stmts.read(); - guard.get(stmt_key).cloned() + let plan = self.prepared_stmts.read().get(stmt_key).cloned(); + metrics::record_cache_lookup("mysql_prepared_stmt", plan.is_some()); + plan } /// Save the prepared statement and return the parameters and result columns @@ -495,7 +511,7 @@ impl AsyncMysqlShim for MysqlInstanceShi raw_query: &'a str, w: StatementMetaWriter<'a, W>, ) -> Result<()> { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let stmt_id = self.prepared_stmts_counter.fetch_add(1, Ordering::Relaxed); let stmt_key = uuid::Uuid::from_u128(stmt_id as u128).to_string(); let (params, columns) = match self @@ -525,7 +541,7 @@ impl AsyncMysqlShim for MysqlInstanceShi ) -> Result<()> { self.session.clear_warnings(); - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_MYSQL_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_MYSQL_BINQUERY, db.as_str()]) @@ -569,7 +585,7 @@ impl AsyncMysqlShim for MysqlInstanceShi query: &'a str, writer: QueryResultWriter<'a, W>, ) -> Result<()> { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_MYSQL_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_MYSQL_TEXTQUERY, db.as_str()]) @@ -1039,6 +1055,34 @@ mod tests { ) } + #[test] + fn test_prepared_statement_cache_metrics() { + use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; + + let hits = CACHE_HIT.with_label_values(&["mysql_prepared_stmt"]); + let misses = CACHE_MISS.with_label_values(&["mysql_prepared_stmt"]); + let before = (hits.get(), misses.get()); + let mut shim = create_shim(); + assert!(shim.plan("stmt").is_none()); + shim.save_plan(SqlPlan::Shortcut("SELECT 1".into()), "stmt".into()) + .unwrap(); + assert!(shim.plan("stmt").is_some()); + shim.do_close("stmt".into()); + assert!(shim.plan("stmt").is_none()); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2)); + } + + #[test] + fn test_batching_context() { + for enabled in [false, true] { + let shim = create_shim().with_batching_enabled(enabled); + let ctx = shim.new_query_context(); + assert_eq!(ctx.batching_enabled(), enabled); + assert!(!ctx.logical_batching_enabled()); + assert_eq!(ctx.channel(), Channel::Mysql); + } + } + fn statement_with_transformed_placeholders(query: &str) -> Statement { let mut statements = ParserContext::create_with_dialect(query, &MySqlDialect {}, ParseOptions::default()) diff --git a/src/servers/src/mysql/server.rs b/src/servers/src/mysql/server.rs index 0f48cb0d598..811fc9827ec 100644 --- a/src/servers/src/mysql/server.rs +++ b/src/servers/src/mysql/server.rs @@ -86,6 +86,7 @@ pub struct MysqlSpawnConfig { reject_no_database: bool, // prepared statement cache capacity prepared_stmt_cache_size: usize, + batching_enabled: bool, } impl MysqlSpawnConfig { @@ -102,9 +103,16 @@ impl MysqlSpawnConfig { keep_alive_secs, reject_no_database, prepared_stmt_cache_size, + batching_enabled: false, } } + /// Enables ordinary-table batching for connections accepted by this server. + pub fn with_batching_enabled(mut self, enabled: bool) -> Self { + self.batching_enabled = enabled; + self + } + fn tls(&self) -> Option> { self.tls.get_config() } @@ -213,7 +221,8 @@ impl MysqlServer { stream.peer_addr()?, process_id, spawn_config.prepared_stmt_cache_size, - ); + ) + .with_batching_enabled(spawn_config.batching_enabled); let (mut r, w) = stream.into_split(); let mut w = BufWriter::with_capacity(DEFAULT_RESULT_SET_WRITE_BUFFER_SIZE, w); diff --git a/src/servers/src/opentsdb/codec.rs b/src/servers/src/opentsdb/codec.rs index c4760aa74de..55674eebf7a 100644 --- a/src/servers/src/opentsdb/codec.rs +++ b/src/servers/src/opentsdb/codec.rs @@ -12,11 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::{Column, ColumnDataType, InsertRequest as GrpcInsertRequest, SemanticType, column}; -use common_query::prelude::{greptime_timestamp, greptime_value}; - -use crate::error::{self, Result}; - #[derive(Debug, Clone)] pub struct DataPoint { metric: String, @@ -35,76 +30,6 @@ impl DataPoint { } } - pub fn try_create(line: &str) -> Result { - let tokens = line.split_whitespace().collect::>(); - let cmd = if tokens.is_empty() { "" } else { tokens[0] }; - // OpenTSDB command is case sensitive, verified in real OpenTSDB. - if cmd != "put" { - return error::InvalidQuerySnafu { - reason: format!("unknown command {cmd}."), - } - .fail(); - } - if tokens.len() < 4 { - return error::InvalidQuerySnafu { - reason: format!( - "put: illegal argument: not enough arguments (need least 4, got {})", - tokens.len() - ), - } - .fail(); - } - - let metric = tokens[1]; - - let ts_millis = match tokens[2].parse::() { - Ok(t) => Self::timestamp_to_millis(t), - Err(_) => { - return error::InvalidQuerySnafu { - reason: format!("put: invalid timestamp: {}", tokens[2]), - } - .fail(); - } - }; - - let value = match tokens[3].parse::() { - Ok(v) => v, - Err(_) => { - return error::InvalidQuerySnafu { - reason: format!("put: invalid value: {}", tokens[3]), - } - .fail(); - } - }; - - let mut tags = Vec::with_capacity(tokens.len() - 4); - for token in tokens.iter().skip(4) { - let tag = token.split('=').collect::>(); - if tag.len() != 2 || tag[0].is_empty() || tag[1].is_empty() { - return error::InvalidQuerySnafu { - reason: format!("put: invalid tag: {token}"), - } - .fail(); - } - let tagk = tag[0].to_string(); - let tagv = tag[1].to_string(); - if tags.iter().any(|(t, _)| t == &tagk) { - return error::InvalidQuerySnafu { - reason: format!("put: illegal argument: duplicate tag: {tagk}"), - } - .fail(); - } - tags.push((tagk, tagv)); - } - - Ok(DataPoint { - metric: metric.to_string(), - ts_millis, - value, - tags, - }) - } - pub fn metric(&self) -> &str { &self.metric } @@ -125,53 +50,6 @@ impl DataPoint { self.value } - pub fn as_grpc_insert(&self) -> GrpcInsertRequest { - let mut columns = Vec::with_capacity(2 + self.tags.len()); - - let ts_column = Column { - column_name: greptime_timestamp().to_string(), - values: Some(column::Values { - timestamp_millisecond_values: vec![self.ts_millis], - ..Default::default() - }), - semantic_type: SemanticType::Timestamp as i32, - datatype: ColumnDataType::TimestampMillisecond as i32, - ..Default::default() - }; - columns.push(ts_column); - - let field_column = Column { - column_name: greptime_value().to_string(), - values: Some(column::Values { - f64_values: vec![self.value], - ..Default::default() - }), - semantic_type: SemanticType::Field as i32, - datatype: ColumnDataType::Float64 as i32, - ..Default::default() - }; - columns.push(field_column); - - for (tagk, tagv) in self.tags.iter() { - columns.push(Column { - column_name: tagk.clone(), - values: Some(column::Values { - string_values: vec![tagv.clone()], - ..Default::default() - }), - semantic_type: SemanticType::Tag as i32, - datatype: ColumnDataType::String as i32, - ..Default::default() - }); - } - - GrpcInsertRequest { - table_name: self.metric.clone(), - columns, - row_count: 1, - } - } - pub fn timestamp_to_millis(t: i64) -> i64 { // 9999999999999 (13 digits) is of date "Sat Nov 20 2286 17:46:39 UTC", // 999999999999 (12 digits) is "Sun Sep 09 2001 01:46:39 UTC", @@ -184,112 +62,3 @@ impl DataPoint { } } } - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_try_create() { - fn test_illegal_line(line: &str, expected_err: &str) { - let result = DataPoint::try_create(line); - match result.unwrap_err() { - error::Error::InvalidQuery { reason, .. } => { - assert_eq!(reason, expected_err) - } - _ => unreachable!(), - } - } - - test_illegal_line("no_put", "unknown command no_put."); - test_illegal_line( - "put", - "put: illegal argument: not enough arguments (need least 4, got 1)", - ); - test_illegal_line( - "put metric.foo notatime 42 host=web01", - "put: invalid timestamp: notatime", - ); - test_illegal_line( - "put metric.foo 1000 notavalue host=web01", - "put: invalid value: notavalue", - ); - test_illegal_line("put metric.foo 1000 42 host=", "put: invalid tag: host="); - test_illegal_line( - "put metric.foo 1000 42 host=web01 host=web02", - "put: illegal argument: duplicate tag: host", - ); - - let data_point = DataPoint::try_create( - "put sys.if.bytes.out 1479496100 1.3E3 host=web01 interface=eth0", - ) - .unwrap(); - assert_eq!(data_point.metric, "sys.if.bytes.out"); - assert_eq!(data_point.ts_millis, 1479496100000); - assert_eq!(data_point.value, 1.3e3); - assert_eq!( - data_point.tags, - vec![ - ("host".to_string(), "web01".to_string()), - ("interface".to_string(), "eth0".to_string()) - ] - ); - - let data_point = - DataPoint::try_create("put sys.procs.running 1479496100 42 host=web01").unwrap(); - assert_eq!(data_point.metric, "sys.procs.running"); - assert_eq!(data_point.ts_millis, 1479496100000); - assert_eq!(data_point.value, 42f64); - assert_eq!( - data_point.tags, - vec![("host".to_string(), "web01".to_string())] - ); - } - - #[test] - fn test_as_grpc_insert() { - let data_point = DataPoint { - metric: "my_metric_1".to_string(), - ts_millis: 1000, - value: 1.0, - tags: vec![ - ("tagk1".to_string(), "tagv1".to_string()), - ("tagk2".to_string(), "tagv2".to_string()), - ], - }; - - let grpc_insert = data_point.as_grpc_insert(); - assert_eq!(grpc_insert.table_name, "my_metric_1"); - - let columns = &grpc_insert.columns; - let row_count = grpc_insert.row_count; - - assert_eq!(row_count, 1); - assert_eq!(columns.len(), 4); - - assert_eq!(columns[0].column_name, greptime_timestamp()); - assert_eq!( - columns[0] - .values - .as_ref() - .unwrap() - .timestamp_millisecond_values, - vec![1000] - ); - - assert_eq!(columns[1].column_name, greptime_value()); - assert_eq!(columns[1].values.as_ref().unwrap().f64_values, vec![1.0]); - - assert_eq!(columns[2].column_name, "tagk1"); - assert_eq!( - columns[2].values.as_ref().unwrap().string_values, - vec!["tagv1"] - ); - - assert_eq!(columns[3].column_name, "tagk2"); - assert_eq!( - columns[3].values.as_ref().unwrap().string_values, - vec!["tagv2"] - ); - } -} diff --git a/src/servers/src/otel_arrow.rs b/src/servers/src/otel_arrow.rs index 0cf07c76468..490ae4b1c58 100644 --- a/src/servers/src/otel_arrow.rs +++ b/src/servers/src/otel_arrow.rs @@ -24,6 +24,7 @@ use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_serv use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ BatchArrowRecords, BatchStatus, StatusCode as ArrowStatusCode, }; +use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::metrics::v1::metric; use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; use tonic::metadata::{Entry, MetadataValue}; @@ -50,23 +51,42 @@ impl OtelArrowServiceHandler { } } +/// Removes unsupported Arrow histograms before ingestion, preserving other metrics. +fn remove_exponential_histograms(request: &mut ExportMetricsServiceRequest) -> bool { + let mut has_data_points = false; + for scope in request + .resource_metrics + .iter_mut() + .flat_map(|resource| &mut resource.scope_metrics) + { + scope.metrics.retain(|item| { + if let Some(metric::Data::ExponentialHistogram(histogram)) = &item.data { + has_data_points |= !histogram.data_points.is_empty(); + false + } else { + true + } + }); + } + has_data_points +} + fn batch_status( batch_id: i64, outcome: MetricsIngestOutcome, has_exponential_histogram_data_points: bool, ) -> BatchStatus { - let status_code = if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { + let status_code = if outcome.accepted_data_points == 0 + && (outcome.rejected_data_points > 0 || has_exponential_histogram_data_points) + { ArrowStatusCode::InvalidArgument } else { ArrowStatusCode::Ok }; - let status_message = match outcome.error_message { - // Arrow keeps the feature gate off, so these fail before per-point validation. - Some(_) if has_exponential_histogram_data_points => { - EXPONENTIAL_HISTOGRAM_UNSUPPORTED.to_string() - } - Some(message) => message, - None => String::new(), + let status_message = if has_exponential_histogram_data_points { + EXPONENTIAL_HISTOGRAM_UNSUPPORTED.to_string() + } else { + outcome.error_message.unwrap_or_default() }; BatchStatus { batch_id, @@ -117,7 +137,7 @@ impl ArrowMetricsService for OtelArrowServiceHandler outcome, Err(error::Error::InvalidOtlpMetricInput { reason }) => { @@ -202,19 +212,61 @@ mod tests { use super::*; #[test] - fn batch_status_explains_arrow_exponential_histogram_limit() { - let status = batch_status( - 7, - MetricsIngestOutcome { - rejected_data_points: 1, - error_message: Some("internal OTLP rejection detail".to_string()), - ..Default::default() - }, - true, - ); + fn removes_arrow_exponential_histograms_and_preserves_other_metrics() { + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, ResourceMetrics, + ScopeMetrics, + }; - assert_eq!(7, status.batch_id); - assert_eq!(ArrowStatusCode::InvalidArgument as i32, status.status_code); - assert_eq!(EXPONENTIAL_HISTOGRAM_UNSUPPORTED, status.status_message); + let mut request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![ + Metric { + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint::default()], + ..Default::default() + })), + ..Default::default() + }, + Metric { + data: Some(metric::Data::Gauge(Gauge::default())), + ..Default::default() + }, + ], + ..Default::default() + }], + ..Default::default() + }], + }; + assert!(remove_exponential_histograms(&mut request)); + let metrics = &request.resource_metrics[0].scope_metrics[0].metrics; + assert_eq!(metrics.len(), 1); + assert!(matches!(metrics[0].data, Some(metric::Data::Gauge(_)))); + assert!(!remove_exponential_histograms(&mut request)); + } + + #[test] + fn batch_status_explains_arrow_exponential_histogram_limit() { + for accepted_data_points in [0, 1] { + let status = batch_status( + 7, + MetricsIngestOutcome { + accepted_data_points, + ..Default::default() + }, + true, + ); + assert_eq!(7, status.batch_id); + assert_eq!( + if accepted_data_points == 0 { + ArrowStatusCode::InvalidArgument + } else { + ArrowStatusCode::Ok + } as i32, + status.status_code + ); + assert_eq!(EXPONENTIAL_HISTOGRAM_UNSUPPORTED, status.status_message); + } } } diff --git a/src/servers/src/otlp/coerce.rs b/src/servers/src/otlp/coerce.rs index 2e7dd36ed7a..2f8d25247cf 100644 --- a/src/servers/src/otlp/coerce.rs +++ b/src/servers/src/otlp/coerce.rs @@ -30,6 +30,9 @@ pub enum TraceCoerceError { // - String to Int64 // - String to Float64 // - String to Boolean +// The empty string and "-" coerce to null instead of failing: instrumentations +// fill numeric attributes from absent headers this way (PSR-7 `getHeaderLine` +// returns ""), and rejecting the value would reject the whole span. // // Lossless signed-to-unsigned integer casts. These let the built-in data // models move from unsigned to signed integers while existing unsigned tables @@ -86,6 +89,13 @@ pub fn coerce_value_data( return Ok(None); }; + if let ValueData::StringValue(s) = v + && matches!(s.as_str(), "" | "-") + && is_supported_trace_coercion(request_type, target) + { + return Ok(None); + } + let Some(value) = coerce_non_null_value(target, request_type, v) else { return Err(TraceCoerceError::Unsupported); }; @@ -273,6 +283,24 @@ mod tests { assert_eq!(result, Err(TraceCoerceError::Unsupported)); } + #[test] + fn test_coerce_absent_value_strings_to_null() { + for s in ["", "-"] { + for target in [ + ColumnDataType::Int64, + ColumnDataType::Float64, + ColumnDataType::Boolean, + ] { + let result = coerce_value_data( + &Some(ValueData::StringValue(s.to_string())), + target, + ColumnDataType::String, + ); + assert_eq!(result, Ok(None), "{s:?} to {target:?}"); + } + } + } + #[test] fn test_coerce_float64_to_int64_not_supported() { let result = coerce_value_data( diff --git a/src/servers/src/otlp/logs.rs b/src/servers/src/otlp/logs.rs index f32816fbe36..53a2189044d 100644 --- a/src/servers/src/otlp/logs.rs +++ b/src/servers/src/otlp/logs.rs @@ -15,7 +15,7 @@ use std::collections::BTreeMap; use ahash::{HashMap, HashMapExt}; -use api::helper::ColumnDataTypeWrapper; +use api::helper::{ColumnDataTypeWrapper, timestamp_unit}; use api::v1::column_data_type_extension::TypeExt; use api::v1::column_def::options_from_column_schema; use api::v1::value::ValueData; @@ -898,16 +898,6 @@ fn is_timestamp_type(datatype: ColumnDataType) -> bool { timestamp_unit(datatype).is_some() } -fn timestamp_unit(datatype: ColumnDataType) -> Option { - match datatype { - ColumnDataType::TimestampSecond => Some(TimeUnit::Second), - ColumnDataType::TimestampMillisecond => Some(TimeUnit::Millisecond), - ColumnDataType::TimestampMicrosecond => Some(TimeUnit::Microsecond), - ColumnDataType::TimestampNanosecond => Some(TimeUnit::Nanosecond), - _ => None, - } -} - fn parse_export_logs_service_request_to_rows( request: ExportLogsServiceRequest, select_info: Box, diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 6b7bbd3fab1..7a485545bdd 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -41,6 +41,7 @@ use table::requests::{ }; use crate::error::{self, Result}; +use crate::metrics::OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; @@ -138,7 +139,7 @@ pub fn to_grpc_insert_requests( && !metric_ctx.is_legacy && let Some(r) = resource.resource.as_ref() { - resource_info.observe(&r.attributes, resource, metric_ctx); + resource_info.observe(&r.attributes, resource); } let resource_attrs = resource.resource.as_ref().map(|r| { @@ -348,12 +349,13 @@ fn from_metric_type(data: &metric::Data) -> MetricType { } } -/// Non-scalar values (bool, arrays, maps, bytes) are not representable as tags. +/// Non-scalar values (arrays, maps, bytes) are not representable as tags. fn scalar_value_string(value: Option<&AnyValue>) -> Option { match value.and_then(|v| v.value.as_ref())? { any_value::Value::StringValue(s) => Some(s.clone()), any_value::Value::IntValue(v) => Some(v.to_string()), any_value::Value::DoubleValue(v) => Some(v.to_string()), + any_value::Value::BoolValue(v) => Some(v.to_string()), _ => None, } } @@ -630,10 +632,13 @@ fn encode_exponential_histogram( metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, ) -> Result { - if let Err(rejection) = exponential_histogram_gate(histogram, metric_ctx) { + if let Err(rejection) = exponential_histogram_gate(histogram) { reject_data_points(outcome, histogram.data_points.len(), || { rejection.message(name) })?; + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS + .with_label_values(&[rejection.reason_label()]) + .inc_by(histogram.data_points.len() as u64); return Ok(false); } @@ -651,6 +656,9 @@ fn encode_exponential_histogram( reject_data_points(outcome, 1, || { format!("metric `{name}` data point {index}: {reason}") })?; + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS + .with_label_values(&["invalid_data_point"]) + .inc(); continue; } }; @@ -684,17 +692,20 @@ fn encode_exponential_histogram( } pub(crate) enum ExponentialHistogramRejection { - Disabled, DeltaTemporality, UnspecifiedTemporality, } impl ExponentialHistogramRejection { + fn reason_label(&self) -> &'static str { + match self { + Self::DeltaTemporality => "delta_temporality", + Self::UnspecifiedTemporality => "unspecified_temporality", + } + } + fn message(&self, name: &str) -> String { match self { - Self::Disabled => format!( - "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" - ), Self::DeltaTemporality => format!( "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" ), @@ -708,13 +719,11 @@ impl ExponentialHistogramRejection { /// Whole-metric acceptance, decided once for the encoder and for the resource /// descriptor, which must not describe a resource whose data was rejected. /// Individual points can still fail [`exponential_histogram_value`]. +/// Raw-delta storage currently supports sums and explicit histograms; exponential +/// histograms still require cumulative input. Delta-to-cumulative conversion is out of scope. pub(crate) fn exponential_histogram_gate( histogram: &ExponentialHistogram, - metric_ctx: &OtlpMetricCtx, ) -> std::result::Result<(), ExponentialHistogramRejection> { - if !metric_ctx.experimental_enable_exponential_histogram { - return Err(ExponentialHistogramRejection::Disabled); - } match AggregationTemporality::try_from(histogram.aggregation_temporality) { Ok(AggregationTemporality::Cumulative) => Ok(()), Ok(AggregationTemporality::Delta) => Err(ExponentialHistogramRejection::DeltaTemporality), @@ -878,25 +887,44 @@ fn convert_bucket_range( merged.push((target_index, count)); } - let length = u32::try_from(merged.len()) - .map_err(|_| format!("{name} bucket span length exceeds u32"))?; - let span = BucketSpan { - offset: merged[0].0, - length, - }; + let mut spans = Vec::::new(); let mut deltas = Vec::with_capacity(merged.len()); + let mut previous_index = None::; let mut previous = 0i64; - for (_, count) in merged { + for (index, count) in merged { + if count == 0 { + continue; + } + match (spans.last_mut(), previous_index) { + (Some(span), Some(prev_index)) if prev_index.checked_add(1) == Some(index) => { + span.length = span + .length + .checked_add(1) + .ok_or_else(|| format!("{name} bucket span length exceeds u32"))?; + } + (_, prev_index) => { + let offset = match prev_index { + Some(prev_index) => index + .checked_sub(prev_index) + .and_then(|gap| gap.checked_sub(1)) + .ok_or_else(|| format!("{name} bucket span offset overflows i32"))?, + None => index, + }; + spans.push(BucketSpan { offset, length: 1 }); + } + } let count = i64::try_from(count) .map_err(|_| format!("{name} bucket count {count} overflows i64"))?; let delta = count .checked_sub(previous) .ok_or_else(|| format!("{name} bucket delta overflows i64"))?; deltas.push(delta); + previous_index = Some(index); + // Bucket deltas continue across span gaps; omitted zero buckets do not reset them. previous = count; } - Ok((vec![span], deltas, total)) + Ok((spans, deltas, total)) } fn downscale_bucket_index(index: i32, downscale_shift: u32) -> std::result::Result { @@ -971,9 +999,15 @@ fn write_timestamp( table: &mut TableData, row: &mut Vec, time_nano: i64, - legacy_mode: bool, + metric_ctx: &OtlpMetricCtx, ) -> Result<()> { - if legacy_mode { + // Keep the full nanosecond precision whenever the request is headed for + // the metric engine: `Inserter::handle_metric_row_inserts` converts the + // timestamps to the physical table's time index unit (which may be + // micro/nanosecond). Only the non-metric prometheus-compatible path is + // fixed to milliseconds, to keep auto-created mito tables on the + // millisecond time index. + if metric_ctx.is_legacy || metric_ctx.with_metric_engine { row_writer::write_ts_to_nanos( table, greptime_timestamp(), @@ -1075,7 +1109,7 @@ fn write_tags_and_timestamp( )?; } - write_timestamp(table, row, timestamp_nanos, metric_ctx.is_legacy)?; + write_timestamp(table, row, timestamp_nanos, metric_ctx)?; Ok(()) } @@ -1117,8 +1151,8 @@ fn encode_gauge( Ok(()) } -/// encode this sum metric -/// +/// Encodes sums, preserving delta points as interval values with a temporality tag. +/// Ingestion is stateless: values are never accumulated across timestamps. fn encode_sum( table_writer: &mut MultiTableData, name: &str, @@ -1178,6 +1212,8 @@ const HISTOGRAM_LE_COLUMN: &str = "le"; /// /// By its Prometheus compatibility, we hope to be able to use prometheus /// quantile functions on this table. +/// Delta points retain their temporality tag and interval values. Bucket counts +/// are prefix-summed within each point, never accumulated across timestamps. fn encode_histogram( table_writer: &mut MultiTableData, name: &str, @@ -1512,6 +1548,7 @@ fn encode_summary( #[cfg(test)] mod tests { + use api::v1::ColumnDataType; use common_query::prelude::set_default_prefix; use otel_arrow_rust::proto::opentelemetry::common::v1::AnyValue; use otel_arrow_rust::proto::opentelemetry::common::v1::any_value::Value as Val; @@ -1535,6 +1572,15 @@ mod tests { } } + fn bool_keyvalue(key: &str, value: bool) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(Val::BoolValue(value)), + }), + } + } + fn descriptor_ctx() -> OtlpMetricCtx { OtlpMetricCtx { resource_info: true, @@ -1698,6 +1744,64 @@ mod tests { assert_eq!(attr_value(&attrs, "instance").as_deref(), Some("inst-1")); } + #[test] + fn test_boolean_attributes_keep_series_distinct() { + set_default_prefix(None).unwrap(); + let request = metrics_request(vec![Metric { + name: "bool_repro".to_string(), + data: Some(metric::Data::Sum(Sum { + data_points: vec![ + NumberDataPoint { + attributes: vec![bool_keyvalue("failed", false)], + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + value: Some(Value::AsDouble(100.0)), + ..Default::default() + }, + NumberDataPoint { + attributes: vec![bool_keyvalue("failed", true)], + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + value: Some(Value::AsDouble(7.0)), + ..Default::default() + }, + ], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + is_monotonic: true, + })), + ..Default::default() + }]); + let conversion = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + let rows = conversion + .requests + .inserts + .iter() + .find(|insert| insert.table_name == "bool_repro_total") + .expect("missing bool_repro_total table") + .rows + .as_ref() + .unwrap(); + assert_eq!(2, rows.rows.len()); + let failed = rows + .schema + .iter() + .position(|column| column.column_name == "failed") + .expect("boolean attribute `failed` was dropped from the label set"); + let values = rows + .rows + .iter() + .map(|row| row.values[failed].value_data.clone()) + .collect::>(); + assert_eq!( + vec![ + Some(ValueData::StringValue("false".to_string())), + Some(ValueData::StringValue("true".to_string())), + ], + values + ); + } + #[test] fn test_encode_gauge() { let mut tables = MultiTableData::default(); @@ -2294,6 +2398,105 @@ mod tests { assert_eq!(total, 21); } + #[test] + fn test_convert_bucket_range_compacts_zero_runs() { + let large = (1u64 << 53) + 1; + for (offset, counts, shift, expected_spans, expected_deltas, expected_total) in [ + ( + -3, + vec![0, 3, 0, 0, 7, 9, 0], + 0, + vec![(-1, 1), (2, 2)], + vec![3, 4, 2], + 19, + ), + ( + -4, + vec![0, 0, 2, 3, 0, 0, 0, 0, 7, 0], + 1, + vec![(0, 1), (2, 1)], + vec![5, 2], + 12, + ), + (0, vec![0, 0, 0], 0, vec![], vec![], 0), + (0, vec![], 0, vec![], vec![], 0), + (i32::MIN, vec![0, 1], 0, vec![(i32::MIN + 2, 1)], vec![1], 1), + (i32::MAX - 1, vec![1], 0, vec![(i32::MAX, 1)], vec![1], 1), + ( + 0, + vec![large, 0, large + 2], + 0, + vec![(1, 1), (1, 1)], + vec![large as i64, 2], + large * 2 + 2, + ), + ] { + let buckets = exponential_buckets(offset, counts); + let (spans, deltas, total) = + convert_bucket_range("positive", Some(&buckets), shift).unwrap(); + assert_eq!( + spans + .iter() + .map(|span| (span.offset, span.length)) + .collect::>(), + expected_spans, + "{buckets:?}, shift={shift}" + ); + assert_eq!(deltas, expected_deltas, "{buckets:?}, shift={shift}"); + assert_eq!(total, expected_total); + } + + // Compaction must not hide invalid source indexes, even for empty buckets. + let buckets = exponential_buckets(i32::MAX, vec![0]); + assert!( + convert_bucket_range("positive", Some(&buckets), 0) + .unwrap_err() + .contains("shifted bucket index overflows") + ); + } + + #[test] + fn test_exponential_histogram_compaction_preserves_integer_counts() { + use common_query::native_histogram::{ + COUNT_I64_FIELD, NEGATIVE_BUCKETS_I64_FIELD, NEGATIVE_SPAN_LENGTHS_FIELD, + NEGATIVE_SPAN_OFFSETS_FIELD, POSITIVE_BUCKETS_I64_FIELD, POSITIVE_SPAN_LENGTHS_FIELD, + POSITIVE_SPAN_OFFSETS_FIELD, + }; + + let large = (1u64 << 53) + 1; + let buckets = exponential_buckets(-2, vec![0, large, 0, 3, 0]); + let point = ExponentialHistogramDataPoint { + count: 2 * (large + 3), + positive: Some(buckets.clone()), + negative: Some(buckets), + ..Default::default() + }; + let (value, _) = exponential_histogram_value(&point).unwrap(); + assert_eq!( + native_field(&value, COUNT_I64_FIELD), + Some(ValueData::I64Value(point.count as i64)) + ); + for (offsets, lengths, counts) in [ + ( + POSITIVE_SPAN_OFFSETS_FIELD, + POSITIVE_SPAN_LENGTHS_FIELD, + POSITIVE_BUCKETS_I64_FIELD, + ), + ( + NEGATIVE_SPAN_OFFSETS_FIELD, + NEGATIVE_SPAN_LENGTHS_FIELD, + NEGATIVE_BUCKETS_I64_FIELD, + ), + ] { + assert_eq!(i32_list(native_field(&value, offsets)), vec![0, 1]); + assert_eq!(i32_list(native_field(&value, lengths)), vec![1, 1]); + assert_eq!( + i64_list(native_field(&value, counts)), + vec![large as i64, 3] + ); + } + } + #[test] fn test_exponential_histogram_value_uses_integer_family() { use common_query::native_histogram::{ @@ -2498,7 +2701,165 @@ mod tests { } #[test] - fn test_exponential_histogram_gate_and_partial_outcome() { + fn test_metric_engine_path_keeps_nanosecond_precision() { + let time_unix_nano = 1_704_067_200_123_456_789u64; + let request = metrics_request(vec![Metric { + name: "my_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }]); + + // The metric engine path keeps the full nanosecond precision here; + // `Inserter::handle_metric_row_inserts` converts the timestamps to + // the physical table's time index unit afterwards. + let mut metric_ctx = OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + }; + let MetricsConversion { requests, .. } = + to_grpc_insert_requests(request, &mut metric_ctx).unwrap(); + + let rows = requests.inserts[0].rows.as_ref().unwrap(); + let ts_index = rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap(); + assert_eq!( + rows.schema[ts_index].datatype, + ColumnDataType::TimestampNanosecond as i32 + ); + assert!(matches!( + rows.rows[0].values[ts_index].value_data, + Some(ValueData::TimestampNanosecondValue( + 1_704_067_200_123_456_789 + )) + )); + + // The non-metric prometheus-compatible path stays millisecond so + // auto-created mito tables keep the millisecond time index. + let mut compat_ctx = OtlpMetricCtx::default(); + let request = metrics_request(vec![Metric { + name: "my_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }]); + let MetricsConversion { requests, .. } = + to_grpc_insert_requests(request, &mut compat_ctx).unwrap(); + + let rows = requests.inserts[0].rows.as_ref().unwrap(); + let ts_index = rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap(); + assert_eq!( + rows.schema[ts_index].datatype, + ColumnDataType::TimestampMillisecond as i32 + ); + assert!(matches!( + rows.rows[0].values[ts_index].value_data, + Some(ValueData::TimestampMillisecondValue(1_704_067_200_123)) + )); + } + + #[test] + fn test_exponential_histogram_rejection_metrics() { + // Other conversion tests update the same process-global counters. + const ISOLATED_ENV: &str = "GREPTIME_TEST_OTLP_REJECTION_METRICS_ISOLATED"; + if std::env::var_os(ISOLATED_ENV).is_none() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "otlp::metrics::tests::test_exponential_histogram_rejection_metrics", + ]) + .env(ISOLATED_ENV, "1") + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed"), + "isolated rejection metrics test failed\nstdout:\n{stdout}\nstderr:\n{}", + String::from_utf8_lossy(&output.stderr) + ); + return; + } + + let mut invalid = exponential_point(); + invalid.scale = -5; + for (temporality, points, reason, accepted, rejected) in [ + ( + AggregationTemporality::Delta, + vec![exponential_point(); 2], + "delta_temporality", + 0, + 2, + ), + ( + AggregationTemporality::Unspecified, + vec![exponential_point(); 2], + "unspecified_temporality", + 0, + 2, + ), + ( + AggregationTemporality::Cumulative, + vec![exponential_point(), invalid], + "invalid_data_point", + 1, + 1, + ), + ( + AggregationTemporality::Cumulative, + vec![exponential_point()], + "invalid_data_point", + 1, + 0, + ), + ] { + let counter = + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS.with_label_values(&[reason]); + let before = counter.get(); + let mut request = + metrics_request(vec![exponential_metric("latency", points, temporality)]); + request.resource_metrics[0].resource = Some(Resource { + attributes: vec![keyvalue("service.name", "api")], + ..Default::default() + }); + let mut ctx = descriptor_ctx(); + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, accepted); + assert_eq!(conversion.outcome.rejected_data_points, rejected); + assert_eq!(counter.get() - before, rejected as u64, "{reason}"); + assert_eq!(conversion.resource_info.is_some(), accepted > 0); + + let before = counter.get(); + let empty = metrics_request(vec![exponential_metric("empty", vec![], temporality)]); + assert_eq!( + to_grpc_insert_requests(empty, &mut ctx) + .unwrap() + .outcome + .rejected_data_points, + 0 + ); + assert_eq!(counter.get(), before, "empty metric: {reason}"); + } + } + + #[test] + fn test_exponential_histogram_default_context_and_partial_outcome() { let request = metrics_request(vec![ Metric { name: "temperature".to_string(), @@ -2512,6 +2873,11 @@ mod tests { vec![exponential_point()], AggregationTemporality::Cumulative, ), + exponential_metric( + "delta_latency", + vec![exponential_point()], + AggregationTemporality::Delta, + ), ]); let MetricsConversion { requests, @@ -2520,20 +2886,20 @@ mod tests { .. } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); - assert_eq!(outcome.accepted_data_points, 1); + assert_eq!(outcome.accepted_data_points, 2); assert_eq!(outcome.rejected_data_points, 1); assert!( outcome .error_message .as_deref() .unwrap() - .contains("otlp.experimental_enable_exponential_histogram") + .contains("only cumulative temporality is supported") ); - assert_eq!(requests.inserts.len(), 1); - assert_eq!(requests.inserts[0].table_name, "temperature"); + assert_eq!(requests.inserts.len(), 2); let semantics = decode(&semantic_index); assert!(semantics.contains_key("temperature")); - assert!(!semantics.contains_key("latency")); + assert!(semantics.contains_key("latency")); + assert!(!semantics.contains_key("delta_latency")); let empty = metrics_request(vec![exponential_metric( "empty", @@ -2566,10 +2932,7 @@ mod tests { AggregationTemporality::Cumulative, ), ]); - let mut ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut ctx = OtlpMetricCtx::default(); let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); assert!( @@ -2589,10 +2952,7 @@ mod tests { ), histogram_metric("latency"), ]); - let mut ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut ctx = OtlpMetricCtx::default(); let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); assert!( @@ -2657,10 +3017,7 @@ mod tests { vec![stale.clone()], temporality, )]); - let mut ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut ctx = OtlpMetricCtx::default(); let MetricsConversion { requests, rows, @@ -2688,15 +3045,11 @@ mod tests { vec![point], AggregationTemporality::Cumulative, )]); - let mut new_ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut new_ctx = OtlpMetricCtx::default(); let new_requests = to_grpc_insert_requests(request.clone(), &mut new_ctx) .unwrap() .requests; let mut legacy_ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, is_legacy: true, ..Default::default() }; @@ -2749,7 +3102,7 @@ mod tests { let request = metrics_request(vec![exponential_metric( "x".repeat(1_000), vec![ExponentialHistogramDataPoint::default()], - AggregationTemporality::Cumulative, + AggregationTemporality::Delta, )]); let outcome = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) .unwrap() diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs index c87e95176c1..4b4f2afcad6 100644 --- a/src/servers/src/otlp/metrics/resource_info.rs +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -30,7 +30,6 @@ use otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue; use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ AggregationTemporality, ResourceMetrics, metric, }; -use session::protocol_ctx::OtlpMetricCtx; use crate::error::Result; use crate::otlp::metrics::{ @@ -90,12 +89,7 @@ pub struct ResourceInfoData { impl ResourceInfoData { /// Takes the raw attributes, before the promote filter runs on them. - pub fn observe( - &mut self, - raw_attrs: &[KeyValue], - resource: &ResourceMetrics, - metric_ctx: &OtlpMetricCtx, - ) { + pub fn observe(&mut self, raw_attrs: &[KeyValue], resource: &ResourceMetrics) { let mut tags = Vec::with_capacity(MAX_PROJECTED_TAGS); let ServiceIdentity { job, instance } = service_identity(raw_attrs); if let Some(job) = job { @@ -119,7 +113,7 @@ impl ResourceInfoData { tags.sort_unstable(); let mut observed: BTreeMap = BTreeMap::new(); - for_each_encoded_time(resource, metric_ctx, |ts| { + for_each_encoded_time(resource, |ts| { let window = ts - ts.rem_euclid(SEMANTIC_GRAPH_WINDOW_NANOS); observed .entry(window) @@ -176,11 +170,7 @@ impl ResourceInfoData { /// Visits the times of the data points the encoder writes rows for, so a /// resource is described exactly where it is measured rather than wherever /// its request happens to reach. -fn for_each_encoded_time( - resource: &ResourceMetrics, - metric_ctx: &OtlpMetricCtx, - mut visit: impl FnMut(i64), -) { +fn for_each_encoded_time(resource: &ResourceMetrics, mut visit: impl FnMut(i64)) { fn visit_all(points: impl Iterator, visit: &mut impl FnMut(i64)) { for ts in points { visit(ts as i64); @@ -214,7 +204,7 @@ fn for_each_encoded_time( visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) } Some(metric::Data::ExponentialHistogram(h)) - if exponential_histogram_gate(h, metric_ctx).is_ok() => + if exponential_histogram_gate(h).is_ok() => { for point in &h.data_points { if let Ok((_, ts)) = exponential_histogram_value(point) { @@ -285,7 +275,7 @@ mod tests { kv("k8s.node.name", "node-a"), kv("os.type", "linux"), ]; - data.observe(&attrs, &gauge_at(&[100, 50]), &OtlpMetricCtx::default()); + data.observe(&attrs, &gauge_at(&[100, 50])); assert_eq!(data.rows.len(), 1); let (tags, windows) = data.rows.iter().next().unwrap(); assert_eq!(windows.values().copied().collect::>(), vec![100]); @@ -298,19 +288,11 @@ mod tests { .all(|(k, _)| k != "os.type" && k != "service.instance.id") ); - data.observe( - &[kv("host.id", "h-2")], - &gauge_at(&[10]), - &OtlpMetricCtx::default(), - ); + data.observe(&[kv("host.id", "h-2")], &gauge_at(&[10])); assert_eq!(data.rows.len(), 2); let mut empty = ResourceInfoData::default(); - empty.observe( - &[kv("os.type", "linux")], - &gauge_at(&[100]), - &OtlpMetricCtx::default(), - ); + empty.observe(&[kv("os.type", "linux")], &gauge_at(&[100])); assert!(empty.into_row_insert_requests().unwrap().is_none()); } @@ -322,7 +304,6 @@ mod tests { data.observe( &[kv("service.name", "api")], &gauge_at(&[window + 1, window + 2, 3 * window + 7]), - &OtlpMetricCtx::default(), ); let windows = data.rows.values().next().unwrap(); @@ -352,20 +333,13 @@ mod tests { }], ..Default::default() }; - let enabled = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; - - for (resource, ctx) in [ - ( - exponential(AggregationTemporality::Cumulative), - OtlpMetricCtx::default(), - ), - (exponential(AggregationTemporality::Delta), enabled), + for temporality in [ + AggregationTemporality::Delta, + AggregationTemporality::Unspecified, ] { + let resource = exponential(temporality); let mut data = ResourceInfoData::default(); - data.observe(&[kv("service.name", "api")], &resource, &ctx); + data.observe(&[kv("service.name", "api")], &resource); assert!(data.into_row_insert_requests().unwrap().is_none()); } } @@ -377,7 +351,6 @@ mod tests { data.observe( &[kv("service.name", "api"), kv("host.id", "h-1")], &gauge_at(&[1_700_000_000_123_456_789]), - &OtlpMetricCtx::default(), ); let requests = data.into_row_insert_requests().unwrap().unwrap(); assert_eq!(requests.inserts.len(), 1); diff --git a/src/servers/src/otlp/metrics/resource_info/tests/delta.rs b/src/servers/src/otlp/metrics/resource_info/tests/delta.rs index 8fa86ae51e1..eb1de565c2f 100644 --- a/src/servers/src/otlp/metrics/resource_info/tests/delta.rs +++ b/src/servers/src/otlp/metrics/resource_info/tests/delta.rs @@ -56,11 +56,7 @@ fn observe_ignores_rejected_classic_histogram_windows() { ..Default::default() }; let mut data = ResourceInfoData::default(); - data.observe( - &[kv("service.name", "api")], - &resource, - &OtlpMetricCtx::default(), - ); + data.observe(&[kv("service.name", "api")], &resource); let windows = data.rows.values().next().unwrap(); assert_eq!( diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index 58ef4fdd7b7..1cd5907157b 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -80,6 +80,7 @@ pub struct PostgresServerHandlerInner { force_tls: bool, param_provider: Arc, + batching_enabled: bool, session: Arc, query_parser: Arc, } @@ -114,7 +115,12 @@ impl PgWireServerHandlers for PostgresServerHandler { } impl MakePostgresServerHandler { - fn make(&self, addr: Option, process_id: u32) -> PostgresServerHandler { + fn make( + &self, + addr: Option, + process_id: u32, + batching_enabled: bool, + ) -> PostgresServerHandler { let session = Arc::new(Session::new( addr, Channel::Postgres, @@ -127,6 +133,7 @@ impl MakePostgresServerHandler { force_tls: self.force_tls, param_provider: self.param_provider.clone(), + batching_enabled, session: session.clone(), query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)), }; diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 8a7b7272447..1fa075a8e17 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -13,7 +13,6 @@ // limitations under the License. use std::fmt::Debug; -use std::sync::Exclusive; use ::auth::{ BEARER_TOKEN_USER, Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, @@ -252,7 +251,7 @@ impl StartupHandler for PostgresServerHandlerInner { auth::save_startup_parameters_to_metadata(client, startup); // check if db is valid - match resolve_db_info(Exclusive::new(client), self.query_handler.clone()).await? { + match resolve_db_info(client, self.query_handler.clone()).await? { DbResolution::Resolved(catalog, schema) => { let metadata = client.metadata_mut(); let _ = metadata.insert(super::METADATA_CATALOG.to_owned(), catalog); @@ -595,13 +594,13 @@ enum DbResolution { /// A function extracted to resolve lifetime and readability issues: async fn resolve_db_info( - client: Exclusive<&mut C>, + client: &mut C, query_handler: ServerSqlQueryHandlerRef, ) -> PgWireResult where C: ClientInfo + Unpin + Send, { - let db_ref = client.into_inner().metadata().get(super::METADATA_DATABASE); + let db_ref = client.metadata().get(super::METADATA_DATABASE); if let Some(db) = db_ref { let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); if query_handler diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 674fe786ac3..9d169811984 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -59,6 +59,14 @@ use crate::postgres::utils::convert_err; use crate::postgres::{PostgresServerHandlerInner, fixtures}; use crate::query_handler::sql::ServerSqlQueryHandlerRef; +impl PostgresServerHandlerInner { + fn new_query_context(&self) -> QueryContextRef { + let mut ctx = self.session.new_query_context(); + Arc::make_mut(&mut ctx).set_batching_enabled(self.batching_enabled); + ctx + } +} + #[async_trait] impl SimpleQueryHandler for PostgresServerHandlerInner { #[tracing::instrument(skip_all, fields(protocol = "postgres"))] @@ -68,7 +76,7 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { C::Error: Debug, PgWireError: From<>::Error>, { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_POSTGRES_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_POSTGRES_SIMPLE_QUERY, db.as_str()]) @@ -82,6 +90,10 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { let parsed_query = self.query_parser.compatibility_parser.parse(query); let query = if let Ok(statements) = &parsed_query { + // Comments, whitespace and empty statements also require EmptyQueryResponse. + if statements.is_empty() { + return Ok(vec![Response::EmptyQuery]); + } statements .iter() .map(statement_to_query) @@ -353,6 +365,9 @@ impl QueryParser for DefaultQueryParser { let parsed_statements = self.compatibility_parser.parse(sql); let (sql, copy_to_stdout_format) = if let Ok(mut statements) = parsed_statements { + if statements.is_empty() { + return Ok(None); + } let first_stmt = statements.remove(0); let format = check_copy_to_stdout(&first_stmt); (statement_to_query(&first_stmt), format) @@ -436,7 +451,7 @@ impl ExtendedQueryHandler for PostgresServerHandlerInner { C::Error: Debug, PgWireError: From<>::Error>, { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_POSTGRES_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_POSTGRES_EXTENDED_QUERY, db.as_str()]) diff --git a/src/servers/src/postgres/server.rs b/src/servers/src/postgres/server.rs index 09bc5015a0b..6e4aa942f9d 100644 --- a/src/servers/src/postgres/server.rs +++ b/src/servers/src/postgres/server.rs @@ -38,6 +38,7 @@ pub struct PostgresServer { make_handler: Arc, tls_server_config: Arc, keep_alive_secs: u64, + batching_enabled: bool, bind_addr: Option, process_manager: Option, } @@ -66,17 +67,25 @@ impl PostgresServer { make_handler, tls_server_config, keep_alive_secs, + batching_enabled: false, bind_addr: None, process_manager, } } + /// Enables ordinary-table batching for connections accepted by this server. + pub fn with_batching_enabled(mut self, enabled: bool) -> Self { + self.batching_enabled = enabled; + self + } + fn accept( &self, io_runtime: Runtime, accepting_stream: AbortableStream, ) -> impl Future + use<> { let handler_maker = self.make_handler.clone(); + let batching_enabled = self.batching_enabled; let tls_server_config = self.tls_server_config.clone(); let process_manager = self.process_manager.clone(); accepting_stream.for_each(move |tcp_stream| { @@ -102,7 +111,8 @@ impl PostgresServer { let _handle = io_runtime.spawn(async move { crate::metrics::METRIC_POSTGRES_CONNECTIONS.inc(); - let pg_handler = Arc::new(handler_maker.make(addr, process_id)); + let pg_handler = + Arc::new(handler_maker.make(addr, process_id, batching_enabled)); let r = process_socket(io_stream, tls_acceptor.clone(), pg_handler).await; crate::metrics::METRIC_POSTGRES_CONNECTIONS.dec(); diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 033a1c665e2..25cb94662c0 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -19,7 +19,7 @@ use std::pin::Pin; use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; -use arrow::array::{Array, AsArray}; +use arrow::array::{Array, ArrayRef, AsArray}; use arrow_pg::encoder::{Encoder, encode_value}; use arrow_pg::list_encoder::encode_list; use arrow_schema::{DataType, TimeUnit}; @@ -37,7 +37,7 @@ use datatypes::json::JsonSettings; use datatypes::prelude::{ConcreteDataType, DataType as _, Value}; use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::{Decimal128Type, IntervalType, TimestampType, jsonb_to_string}; -use datatypes::value::StructValue; +use datatypes::value::try_value_from_array; use futures::Stream; use pg_interval::Interval as PgInterval; use pgwire::api::Type; @@ -234,7 +234,8 @@ impl ToSqlText for OidAliasValue<'_> { } } -/// this function will encode greptime's `StructValue` into PostgreSQL jsonb type +/// this function will encode greptime's structured values (`StructValue`, or a +/// list of them) into PostgreSQL jsonb type /// /// Note that greptimedb has different types of StructValue for storing json data, /// based on policy defined in `JsonSettings`. But here the `StructValue` @@ -245,18 +246,35 @@ impl ToSqlText for OidAliasValue<'_> { /// arrays: element in array must be the same type fn encode_struct( _query_ctx: &QueryContextRef, - struct_value: StructValue, + value: Value, builder: &mut S, pg_field: &FieldInfo, ) -> PgWireResult<()> { let encoding_setting = JsonSettings::default(); let json_value = encoding_setting - .decode(Value::Struct(struct_value)) + .decode(value) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; builder.encode_field(&json_value, pg_field) } +/// Encodes one row of a structured column (`Struct`, or `List` of `Struct`) as +/// a PostgreSQL `json` value. +/// +/// Returns an error for arrow field types greptimedb cannot represent +/// (e.g. `Decimal256`) instead of panicking, so the client receives a proper +/// error response. +fn encode_structured_json( + query_ctx: &QueryContextRef, + column: &ArrayRef, + i: usize, + encoder: &mut S, + pg_field: &FieldInfo, +) -> PgWireResult<()> { + let value = try_value_from_array(column.as_ref(), i).map_err(convert_err)?; + encode_struct(query_ctx, value, encoder, pg_field) +} + pub(crate) struct RecordBatchRowStream where S: Encoder, @@ -369,6 +387,11 @@ where } } + DataType::List(list_field) + if matches!(list_field.data_type(), DataType::Struct(_)) => + { + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; + } DataType::List(_) => { let array = column.as_list::(); let items = array.value(i); @@ -376,7 +399,7 @@ where encode_list(encoder, items, pg_field)?; } DataType::Struct(_) => { - encode_struct(query_ctx, Default::default(), encoder, pg_field)?; + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; } DataType::Utf8 => { let arrow_field = arrow_schema.field(j); @@ -445,7 +468,9 @@ pub(super) fn type_gt_to_pg(origin: &ConcreteDataType) -> Result { &ConcreteDataType::Decimal128(_) => Ok(Type::NUMERIC_ARRAY), &ConcreteDataType::Json(_) => Ok(Type::JSON_ARRAY), &ConcreteDataType::Duration(_) => Ok(Type::INTERVAL_ARRAY), - &ConcreteDataType::Struct(_) => Ok(Type::JSON_ARRAY), + // a list of structs is encoded as a single JSON array, matching the + // struct-to-json mapping above + &ConcreteDataType::Struct(_) => Ok(Type::JSON), &ConcreteDataType::Dictionary(_) | &ConcreteDataType::Vector(_) | &ConcreteDataType::List(_) => server_error::UnsupportedDataTypeSnafu { @@ -464,45 +489,6 @@ pub(super) fn type_gt_to_pg(origin: &ConcreteDataType) -> Result { } } -#[allow(dead_code)] -pub(super) fn type_pg_to_gt(origin: &Type) -> Result { - // Note that we only support a small amount of pg data types - match origin { - &Type::BOOL => Ok(ConcreteDataType::boolean_datatype()), - &Type::INT2 => Ok(ConcreteDataType::int16_datatype()), - &Type::INT4 => Ok(ConcreteDataType::int32_datatype()), - &Type::INT8 => Ok(ConcreteDataType::int64_datatype()), - &Type::NUMERIC => Ok(ConcreteDataType::uint64_datatype()), - &Type::VARCHAR | &Type::CHAR | &Type::TEXT => Ok(ConcreteDataType::string_datatype()), - &Type::TIMESTAMP | &Type::TIMESTAMPTZ => Ok(ConcreteDataType::timestamp_datatype( - common_time::timestamp::TimeUnit::Millisecond, - )), - &Type::DATE => Ok(ConcreteDataType::date_datatype()), - &Type::TIME => Ok(ConcreteDataType::timestamp_datatype( - common_time::timestamp::TimeUnit::Microsecond, - )), - &Type::INT2_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::int16_datatype(), - ))), - &Type::INT4_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::int32_datatype(), - ))), - &Type::INT8_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::int64_datatype(), - ))), - &Type::NUMERIC_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::uint64_datatype(), - ))), - &Type::VARCHAR_ARRAY | &Type::CHAR_ARRAY | &Type::TEXT_ARRAY => Ok( - ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::string_datatype())), - ), - _ => server_error::InternalSnafu { - err_msg: format!("unimplemented datatype {origin:?}"), - } - .fail(), - } -} - pub(super) fn invalid_parameter_error(msg: &str, detail: Option) -> PgWireError { let mut error_info = PgErrorCode::Ec22023.to_err_info(msg.to_string()); error_info.detail = detail; @@ -1921,6 +1907,282 @@ mod test { } } + #[test] + fn test_encode_list_of_struct_as_json() { + use datatypes::types::{StructField, StructType}; + use datatypes::value::{ListValue, StructValue}; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + // struct type: {a: int32, b: float64} + let fields = vec![ + StructField::new("a".to_string(), ConcreteDataType::int32_datatype(), true), + StructField::new("b".to_string(), ConcreteDataType::float64_datatype(), true), + ]; + let struct_type = + ConcreteDataType::struct_datatype(StructType::new(Arc::new(fields.clone()))); + let list_type = ConcreteDataType::list_datatype(Arc::new(struct_type.clone())); + + let make_struct = |a: i32, b: f64| { + Value::Struct(StructValue::new( + vec![Value::Int32(a), Value::Float64(b.into())], + StructType::new(Arc::new(fields.clone())), + )) + }; + let rows = vec![ + Value::List(ListValue::new( + vec![make_struct(1, 1.5), Value::Null, make_struct(3, 3.5)], + Arc::new(struct_type.clone()), + )), + Value::List(ListValue::new(vec![], Arc::new(struct_type.clone()))), + Value::Null, + ]; + let vector = Helper::try_from_row_into_vector(&rows, &list_type).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "l".to_string(), + list_type.clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + // The list-of-struct column is declared as a single json value. + let pg_schema = Arc::new(vec![FieldInfo::new( + "l".into(), + None, + None, + type_gt_to_pg(&list_type).unwrap(), + FieldFormat::Text, + )]); + assert_eq!(pg_schema[0].datatype(), &Type::JSON); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL list row stays NULL. + assert_eq!(decoded[2], None); + + // A list of structs is a JSON array; null elements stay null. + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!([ + {"a": 1, "b": 1.5}, + null, + {"a": 3, "b": 3.5}, + ]) + ); + + // An empty list is an empty JSON array. + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!(json, serde_json::json!([])); + } + + #[test] + fn test_encode_struct_value_as_json() { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use api::helper::pb_value_to_value_ref; + use api::v1::Value as PbValue; + use common_query::native_histogram::{ + encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, + }; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + fn histogram_value(histogram: &Histogram) -> Value { + let value_data = encode_native_histogram(histogram).unwrap(); + let pb_value = PbValue { + value_data: Some(value_data), + }; + let column_schema = native_histogram_column_schema().unwrap(); + Value::from(pb_value_to_value_ref( + &pb_value, + column_schema.datatype_extension.as_ref(), + )) + } + + // One integer-count histogram and one float-count histogram, mirroring + // what prometheus remote-write v2 stores. The expected JSON matches the + // HTTP query output for the same payloads. + let int_histogram = Histogram { + schema: 1, + sum: 10.0, + zero_threshold: 0.001, + reset_hint: 2, + start_timestamp: 1500, + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + count: Some(Count::CountInt(8)), + zero_count: Some(ZeroCount::ZeroCountInt(1)), + ..Default::default() + }; + let float_histogram = Histogram { + schema: 2, + sum: 20.0, + zero_threshold: 0.002, + reset_hint: 3, + start_timestamp: 2500, + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + count: Some(Count::CountFloat(6.0)), + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + ..Default::default() + }; + + let rows = vec![ + histogram_value(&int_histogram), + histogram_value(&float_histogram), + Value::Null, + ]; + let vector = + Helper::try_from_row_into_vector(&rows, native_histogram_value_type()).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "greptime_native_histogram", + native_histogram_value_type().clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + let pg_schema = Arc::new(vec![FieldInfo::new( + "greptime_native_histogram".into(), + None, + None, + Type::JSON, + FieldFormat::Text, + )]); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL struct row stays NULL. + assert_eq!(decoded[2], None); + + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + } + #[test] fn test_encode_utf8_oid_alias_data() { let aliases = [ diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index c9aabfb2745..06729795b5b 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -27,6 +27,7 @@ pub(crate) mod v2; pub mod validation; use bytes::Bytes; +use common_memory_manager::MemoryGuard; use lazy_static::lazy_static; use object_pool::Pool; use snafu::ResultExt; @@ -35,7 +36,12 @@ use crate::error; use crate::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest}; use crate::prom_remote_write::row_builder::TablesBuilder; use crate::prom_remote_write::validation::PromValidationMode; -use crate::prom_store::{snappy_decompress, zstd_decompress}; +use crate::prom_store::{ + ChargedBuffer, MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited, + zstd_decompress_limited, +}; +use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; /// Prometheus remote write protocol versions, also used as the `version` label /// of the remote write metrics. @@ -47,43 +53,74 @@ lazy_static! { Pool::new(256, PromWriteRequest::default); } -pub fn try_decompress(is_zstd: bool, body: &[u8]) -> crate::error::Result> { - if is_zstd { - zstd_decompress(body) - } else { - snappy_decompress(body) +/// Decompresses a remote-write body with the codec selected by `is_zstd`, +/// bounded by the hard decoded-size cap and charged to the aggregate +/// request-memory limiter. +/// +/// Due to vmagent's limitation there is a chance that vmagent sends the +/// content-encoding header wrong, so decoding falls back to the other codec +/// when the first attempt fails on malformed input. Size-limit and quota +/// errors are final: the body either expanded beyond the cap under the +/// declared codec, or it could not be admitted, and the other codec cannot +/// produce a different outcome. +/// +/// see +/// see +pub(crate) async fn decompress_remote_write_body( + is_zstd: bool, + body: &[u8], + limiter: &ServerMemoryLimiter, +) -> crate::error::Result { + let decompress = |is_zstd: bool| async move { + if is_zstd { + zstd_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + } else { + snappy_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + } + }; + + match decompress(is_zstd).await { + Ok(buf) => Ok(buf), + Err(e) => { + if matches!( + e, + crate::error::Error::DecompressSnappyPromRemoteRequest { .. } + | crate::error::Error::DecompressZstdPromRemoteRequest { .. } + ) { + decompress(!is_zstd).await + } else { + Err(e) + } + } } } -pub fn decode_remote_write_request( +pub async fn decode_remote_write_request( is_zstd: bool, body: Bytes, prom_validation_mode: PromValidationMode, processor: &mut PromSeriesProcessor, -) -> crate::error::Result> { + limiter: &ServerMemoryLimiter, +) -> crate::error::Result<( + TablesBuilder<'static>, + Vec>, +)> { let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V1_VERSION]) .start_timer(); - // due to vmagent's limitation, there is a chance that vmagent is - // sending content type wrong so we have to apply a fallback with decoding - // the content in another method. - // - // see https://github.com/VictoriaMetrics/VictoriaMetrics/issues/5301 - // see https://github.com/GreptimeTeam/greptimedb/issues/3929 - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - // fallback to the other compression method - try_decompress(!is_zstd, &body[..])? - }; + // Holds the memory permits for the decompressed bytes; the caller must + // keep them alive as long as the returned `TablesBuilder`, which retains + // the decompressed buffer as its raw data. + let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; // Decompression copied the payload out, so the compressed body is no longer needed. drop(body); - + let (data, guards) = buf.into_parts(); let mut request = PROM_WRITE_REQUEST_POOL.pull(PromWriteRequest::default); request - .decode(buf, prom_validation_mode, processor) + .decode(data, prom_validation_mode, processor) .context(error::DecodePromRemoteRequestSnafu)?; - Ok(std::mem::take(&mut request.table_data)) + let tables = std::mem::take(&mut request.table_data); + Ok((tables, guards)) } diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 2f1f4d410c4..9178748bc77 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -50,12 +50,13 @@ use table::requests::{ use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; use crate::prom_remote_write::validation::validate_label_name; -use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; +use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, decompress_remote_write_body}; #[allow(deprecated)] use crate::prom_store::{ DATABASE_LABEL, DATABASE_LABEL_ALT, METRIC_NAME_LABEL, PHYSICAL_TABLE_LABEL, PHYSICAL_TABLE_LABEL_ALT, SCHEMA_LABEL, }; +use crate::request_memory_limiter::ServerMemoryLimiter; use crate::row_writer::{self, TableData}; use crate::semantic::{ METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_GAUGE_HISTOGRAM, METRIC_TYPE_HISTOGRAM, @@ -127,22 +128,18 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub semantic_index: SemanticIndexes, } -pub(crate) fn decode_remote_write_v2( +pub(crate) async fn decode_remote_write_v2( is_zstd: bool, body: Bytes, - native_histograms_enabled: bool, + limiter: &ServerMemoryLimiter, ) -> Result { let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V2_VERSION]) .start_timer(); - // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a - // mismatched content-encoding header, so try the other compression on failure. - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; + // Holds the memory permits for the decompressed bytes until the protobuf + // decoding and conversion below are finished. + let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; // Decompression copied the payload out, so the compressed body is no longer needed. drop(body); let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; @@ -151,13 +148,10 @@ pub(crate) fn decode_remote_write_v2( let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["convert", REMOTE_WRITE_V2_VERSION]) .start_timer(); - convert_remote_write_v2(request, native_histograms_enabled) + convert_remote_write_v2(request) } -fn convert_remote_write_v2( - request: BorrowedRequest<'_>, - native_histograms_enabled: bool, -) -> Result { +fn convert_remote_write_v2(request: BorrowedRequest<'_>) -> Result { ensure!( request.symbols.first().copied() == Some(""), error::InvalidPromRemoteRequestSnafu { @@ -179,14 +173,6 @@ fn convert_remote_write_v2( let counts = scan_series(series, &mut labels_refs, &mut metadata) .context(error::DecodePromRemoteRequestSnafu)?; - ensure!( - native_histograms_enabled || counts.histograms == 0, - error::InvalidPromRemoteRequestSnafu { - msg: "prometheus remote write v2 native histogram ingestion is experimental; set prom_store.experimental_enable_prometheus_native_histogram = true to enable it" - .to_string(), - } - ); - if counts.samples == 0 && counts.histograms == 0 { decode_series_leaves(series, None, Vec::new(), 0, &mut scratch)?; continue; @@ -672,7 +658,7 @@ fn resolve_series_labels<'a>( let mut tags = Vec::with_capacity(labels_refs.len() / 2); label_names.clear(); - for pair in labels_refs.chunks_exact(2) { + for pair in labels_refs.as_chunks::<2>().0 { let name = symbol_ref(symbols, pair[0], "label name")?; let value = symbol_ref(symbols, pair[1], "label value")?; validate_label(name)?; @@ -801,8 +787,9 @@ pub mod test_util { use snafu::ResultExt; use crate::error::{self, Result}; - use crate::prom_remote_write::try_decompress; + use crate::prom_remote_write::decompress_remote_write_body; use crate::prom_store::snappy_compress; + use crate::request_memory_limiter::ServerMemoryLimiter; pub fn request_with_labels_and_samples( labels: Vec<(&str, &str)>, @@ -819,27 +806,44 @@ pub mod test_util { } pub fn decode_request(is_zstd: bool, body: Bytes) -> Result { - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; + let buf = block_on_decode(is_zstd, body)?; Request::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } + /// Runs the charged decompression on a throwaway runtime so plain `#[test]` + /// callers can stay synchronous. + fn block_on_decode(is_zstd: bool, body: Bytes) -> Result> { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async move { + decompress_remote_write_body(is_zstd, &body[..], &ServerMemoryLimiter::default()) + .await + .map(|buf| buf.to_vec()) + }) + } + pub fn write_requests( request: Request, ) -> Result<(Vec, Vec, u64, u64)> { let body = Bytes::from(snappy_compress(&request.encode_to_vec())?); - decode_write_requests(false, body, true) + decode_write_requests(false, body) } pub fn decode_write_requests( is_zstd: bool, body: Bytes, - native_histograms_enabled: bool, ) -> Result<(Vec, Vec, u64, u64)> { - let requests = super::decode_remote_write_v2(is_zstd, body, native_histograms_enabled)?; + let requests = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(super::decode_remote_write_v2( + is_zstd, + body, + &ServerMemoryLimiter::default(), + ))?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -850,11 +854,10 @@ pub mod test_util { pub fn decode_uncompressed_write_requests( body: &[u8], - native_histograms_enabled: bool, ) -> Result<(Vec, Vec, u64, u64)> { let request = super::BorrowedRequest::decode(body).context(error::DecodePromRemoteRequestSnafu)?; - let requests = super::convert_remote_write_v2(request, native_histograms_enabled)?; + let requests = super::convert_remote_write_v2(request)?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -954,9 +957,7 @@ mod tests { assert_eq!(decoded.timeseries[0].samples[0].value, 42.0); assert_eq!(decoded.timeseries[0].metadata.as_ref().unwrap().r#type, 1); assert_eq!( - decode_remote_write_v2(true, body, true) - .unwrap() - .sample_count, + decode_v2_on_test_runtime(true, body).unwrap().sample_count, 1 ); } @@ -990,7 +991,7 @@ mod tests { wire.extend(encoded_message_field(5, &packed_u32_field(1, &[99]))); wire.extend(string_field(4, b"http_requests_total")); - let requests = decode_wire(&wire, true).unwrap(); + let requests = decode_wire(&wire).unwrap(); assert_eq!(requests.sample_count, 2); assert_eq!(requests.histogram_count, 0); let rows = requests.samples.all_req().next().unwrap().rows.unwrap(); @@ -1028,7 +1029,7 @@ mod tests { series.extend(packed_u32_field(1, &[1, 2])); let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); - let requests = decode_wire(&wire, true).unwrap(); + let requests = decode_wire(&wire).unwrap(); assert_eq!(requests.histogram_count, 1); let rows = requests.histograms.all_req().next().unwrap().rows.unwrap(); assert_eq!( @@ -1093,7 +1094,7 @@ mod tests { ("invalid sample", invalid_sample), ("invalid histogram", invalid_histogram), ] { - let error = decode_wire_error(&wire, true, name); + let error = decode_wire_error(&wire, name); assert!( matches!(error, error::Error::DecodePromRemoteRequest { .. }), "{name}: {error}" @@ -1108,7 +1109,7 @@ mod tests { series.extend(encoded_message_field(tag, &[0x08])); let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); - let error = decode_wire_error(&wire, true, "malformed ignored message"); + let error = decode_wire_error(&wire, "malformed ignored message"); assert!(matches!( error, error::Error::DecodePromRemoteRequest { .. } @@ -1174,8 +1175,8 @@ mod tests { }) .unwrap(); assert_eq!(requests.sample_count, 0); - assert!(decode_wire(&[], true).is_err()); - assert!(decode_wire(&[0x0a, 0x00], true).is_err()); + assert!(decode_wire(&[]).is_err()); + assert!(decode_wire(&[0x0a, 0x00]).is_err()); } #[test] @@ -1205,22 +1206,34 @@ mod tests { } #[test] - fn test_fused_decoder_pins_experimental_error_precedence() { - let histogram = series_wire(&[1, 2], 3, &Histogram::default().encode_to_vec()); - let malformed_sample = series_wire(&[1, 2], 2, &[0x08]); + fn test_fused_decoder_rejects_malformed_series_with_histograms() { + let histogram = series_wire( + &[1, 2], + 3, + &Histogram { + count: Some(Count::CountInt(0)), + zero_count: Some(ZeroCount::ZeroCountInt(0)), + ..Default::default() + } + .encode_to_vec(), + ); + let malformed_sample = series_wire(&[1, 3], 2, &[0x08]); let wire = request_wire( - &["", METRIC_NAME_LABEL, "metric"], + &["", METRIC_NAME_LABEL, "histogram", "sample"], &[histogram.clone(), malformed_sample.clone()], ); - let error = decode_wire_error(&wire, false, "histogram before malformed series"); - assert!(error.to_string().contains("ingestion is experimental")); + let error = decode_wire_error(&wire, "histogram before malformed series"); + assert!(matches!( + error, + error::Error::DecodePromRemoteRequest { .. } + )); let wire = request_wire( - &["", METRIC_NAME_LABEL, "metric"], + &["", METRIC_NAME_LABEL, "histogram", "sample"], &[malformed_sample, histogram.clone()], ); - let error = decode_wire_error(&wire, false, "malformed series before histogram"); + let error = decode_wire_error(&wire, "malformed series before histogram"); assert!(matches!( error, error::Error::DecodePromRemoteRequest { .. } @@ -1231,7 +1244,7 @@ mod tests { &["", METRIC_NAME_LABEL, "metric", "job", "api"], &[missing_name, histogram], ); - let error = decode_wire_error(&wire, false, "conversion error before histogram"); + let error = decode_wire_error(&wire, "conversion error before histogram"); assert!(error.to_string().contains("missing '__name__'")); } @@ -2026,16 +2039,27 @@ mod tests { )); } - fn decode_wire( - wire: &[u8], - native_histograms_enabled: bool, - ) -> Result { + fn decode_wire(wire: &[u8]) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(wire).unwrap()); - decode_remote_write_v2(false, body, native_histograms_enabled) + decode_v2_on_test_runtime(false, body) } - fn decode_wire_error(wire: &[u8], native_histograms_enabled: bool, name: &str) -> error::Error { - match decode_wire(wire, native_histograms_enabled) { + /// Runs the async (charged) v2 decoder on a throwaway runtime so plain + /// `#[test]` callers can stay synchronous. + fn decode_v2_on_test_runtime(is_zstd: bool, body: Bytes) -> Result { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(decode_remote_write_v2( + is_zstd, + body, + &ServerMemoryLimiter::default(), + )) + } + + fn decode_wire_error(wire: &[u8], name: &str) -> error::Error { + match decode_wire(wire) { Ok(_) => panic!("{name}: expected decoder error"), Err(error) => error, } @@ -2104,16 +2128,9 @@ mod tests { } fn decode_test_request(request: Request) -> Result { - decode_test_request_with_histograms(request, true) - } - - fn decode_test_request_with_histograms( - request: Request, - native_histograms_enabled: bool, - ) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); - decode_remote_write_v2(false, body, native_histograms_enabled) + decode_v2_on_test_runtime(false, body) } fn assert_invalid(name: &str, request: Request, expected: &str) { diff --git a/src/servers/src/prom_row_builder.rs b/src/servers/src/prom_row_builder.rs index 8b4e35bcd32..43b3a464857 100644 --- a/src/servers/src/prom_row_builder.rs +++ b/src/servers/src/prom_row_builder.rs @@ -23,14 +23,17 @@ use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Rows, SemanticType}; use arrow::array::{ - ArrayRef, Float64Builder, StringBuilder, TimestampMicrosecondBuilder, - TimestampMillisecondBuilder, TimestampNanosecondBuilder, TimestampSecondBuilder, - new_null_array, + ArrayRef, ArrowPrimitiveType, Float64Builder, PrimitiveBuilder, StringBuilder, new_null_array, +}; +use arrow::datatypes::{ + DataType as ArrowDataType, Schema as ArrowSchema, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; -use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use arrow_schema::TimeUnit; use common_query::prelude::{greptime_timestamp, greptime_value}; +use common_time::Timestamp; +use common_time::timestamp::TimeUnit as CommonTimeUnit; use datatypes::data_type::DataType; use datatypes::prelude::ConcreteDataType; use snafu::{OptionExt, ResultExt, ensure}; @@ -129,16 +132,7 @@ pub(crate) fn rows_to_aligned_record_batch( ArrowDataType::Float64 => { source_map.insert(&target_field_name, (src_idx, src_arrow_type)); } - ArrowDataType::Timestamp(unit, _) => { - ensure!( - unit == &TimeUnit::Millisecond, - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Unexpected remote write batch timestamp unit, expect millisecond, got: {}", - unit - ) - } - ); + ArrowDataType::Timestamp(_, _) => { source_map.insert(&target_ts_name, (src_idx, src_arrow_type)); } ArrowDataType::Utf8 => { @@ -156,15 +150,28 @@ pub(crate) fn rows_to_aligned_record_batch( } } - // Build columns in target schema order + // Build columns in target schema order. The timestamp column is built in + // the TARGET schema's unit; `build_arrow_array` converts the request's + // encoding unit (flooring on narrowing) into it. let mut columns = Vec::with_capacity(target_schema.fields().len()); for target_field in target_schema.fields() { if let Some((src_idx, src_arrow_type)) = source_map.get(target_field.name().as_str()) { + let target_type = if matches!( + (src_arrow_type, target_field.data_type()), + ( + ArrowDataType::Timestamp(_, _), + ArrowDataType::Timestamp(_, _) + ) + ) { + target_field.data_type().clone() + } else { + src_arrow_type.clone() + }; let array = build_arrow_array( rows, *src_idx, &rows.schema[*src_idx].column_name, - src_arrow_type.clone(), + target_type, row_count, )?; columns.push(array); @@ -199,6 +206,65 @@ pub(crate) fn identify_missing_columns_from_proto( Ok(missing) } +/// Builds a timestamp array of `T`'s unit from a proto column, appending +/// each value directly into the builder. Values already in the target unit +/// are appended without conversion (the unchanged-unit fast path); others +/// are converted via `Timestamp::convert_to`, flooring on narrowing. +fn build_timestamp_array>( + rows: &Rows, + col_idx: usize, + column_name: &str, + row_count: usize, + target_unit: CommonTimeUnit, +) -> Result { + let mut builder = PrimitiveBuilder::::with_capacity(row_count); + for row in &rows.rows { + let Some(value) = row.values[col_idx].value_data.as_ref() else { + builder.append_null(); + continue; + }; + let (source_unit, raw) = match value { + ValueData::TimestampSecondValue(v) => (CommonTimeUnit::Second, *v), + ValueData::TimestampMillisecondValue(v) => (CommonTimeUnit::Millisecond, *v), + ValueData::DatetimeValue(v) | ValueData::TimestampMicrosecondValue(v) => { + (CommonTimeUnit::Microsecond, *v) + } + ValueData::TimestampNanosecondValue(v) => (CommonTimeUnit::Nanosecond, *v), + v => { + return error::InvalidPromRemoteRequestSnafu { + msg: format!("Unexpected value: {:?}", v), + } + .fail(); + } + }; + if source_unit == target_unit { + builder.append_value(raw); + } else { + let timestamp = Timestamp::new(raw, source_unit); + let Some(converted) = timestamp.convert_to(target_unit) else { + return error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Timestamp value in column '{column_name}' overflows when converting to unit {target_unit:?}" + ), + } + .fail(); + }; + builder.append_value(converted.value()); + } + } + Ok(Arc::new(builder.finish()) as ArrayRef) +} + +/// Converts an arrow time unit to the common time unit. +fn arrow_time_unit(unit: TimeUnit) -> CommonTimeUnit { + match unit { + TimeUnit::Second => CommonTimeUnit::Second, + TimeUnit::Millisecond => CommonTimeUnit::Millisecond, + TimeUnit::Microsecond => CommonTimeUnit::Microsecond, + TimeUnit::Nanosecond => CommonTimeUnit::Nanosecond, + } +} + /// Build a `Vec` suitable for creating a new Prometheus logical table /// directly from the proto `rows.schema`, avoiding the round-trip through Arrow schema. pub fn build_prom_create_table_schema_from_proto( @@ -207,13 +273,22 @@ pub fn build_prom_create_table_schema_from_proto( rows_schema .iter() .map(|col| { - let semantic_type = if col.datatype == api::v1::ColumnDataType::TimestampMillisecond as i32 { + let datatype = api::v1::ColumnDataType::try_from(col.datatype).map_err(|_| { + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Failed to build create table schema, column '{}' has unknown datatype {}", + col.column_name, col.datatype + ), + } + .build() + })?; + let semantic_type = if api::helper::timestamp_unit(datatype).is_some() { SemanticType::Timestamp - } else if col.datatype == api::v1::ColumnDataType::Float64 as i32 { + } else if datatype == api::v1::ColumnDataType::Float64 { SemanticType::Field } else { // tag columns must be String type - ensure!(col.datatype == api::v1::ColumnDataType::String as i32, error::InvalidPromRemoteRequestSnafu{ + ensure!(datatype == api::v1::ColumnDataType::String, error::InvalidPromRemoteRequestSnafu{ msg: format!( "Failed to build create table schema, tag column '{}' must be String but got datatype {}", col.column_name, col.datatype @@ -268,25 +343,45 @@ fn build_arrow_array( StringBuilder::with_capacity(row_count, 0), ValueData::StringValue(v) => v ), - arrow::datatypes::DataType::Timestamp(u, _) => match u { - TimeUnit::Second => build_array!( - TimestampSecondBuilder::with_capacity(row_count), - ValueData::TimestampSecondValue(v) => *v - ), - TimeUnit::Millisecond => build_array!( - TimestampMillisecondBuilder::with_capacity(row_count), - ValueData::TimestampMillisecondValue(v) => *v - ), - TimeUnit::Microsecond => build_array!( - TimestampMicrosecondBuilder::with_capacity(row_count), - ValueData::DatetimeValue(v) => *v, - ValueData::TimestampMicrosecondValue(v) => *v - ), - TimeUnit::Nanosecond => build_array!( - TimestampNanosecondBuilder::with_capacity(row_count), - ValueData::TimestampNanosecondValue(v) => *v - ), - }, + arrow::datatypes::DataType::Timestamp(u, _) => { + // Accept any timestamp encoding and append directly into the + // target-unit builder — no intermediate column allocation. + // Values already in the target unit (the common unchanged + // millisecond case) are appended as-is; others are converted + // first, flooring on narrowing (same semantics as + // `Timestamp::convert_to` on the ordinary insert path). + let target_unit = arrow_time_unit(u); + match u { + TimeUnit::Second => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + TimeUnit::Millisecond => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + TimeUnit::Microsecond => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + TimeUnit::Nanosecond => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + } + } ty => { return error::InvalidPromRemoteRequestSnafu { msg: format!( @@ -305,7 +400,9 @@ fn build_arrow_array( mod tests { use api::v1::value::ValueData; use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value}; - use arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; + use arrow::array::{ + Array, Float64Array, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray, + }; use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use super::{ @@ -313,6 +410,105 @@ mod tests { rows_to_aligned_record_batch, }; + #[test] + fn test_rows_to_aligned_record_batch_converts_time_units() { + // Microsecond-encoded rows aligned to a millisecond schema: narrowing + // floors towards negative infinity (-1001us -> -2ms), matching + // `Timestamp::convert_to` on the ordinary insert path. + let rows = Rows { + schema: vec![ + ColumnSchema { + column_name: "greptime_timestamp".to_string(), + datatype: ColumnDataType::TimestampMicrosecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "greptime_value".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }, + ], + rows: vec![ + Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMicrosecondValue(1000)), + }, + Value { + value_data: Some(ValueData::F64Value(1.0)), + }, + ], + }, + Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMicrosecondValue(-1001)), + }, + Value { + value_data: Some(ValueData::F64Value(2.0)), + }, + ], + }, + ], + }; + let target = ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("greptime_value", DataType::Float64, true), + ]); + + let (batch, _) = rows_to_aligned_record_batch(&rows, &target) + .unwrap() + .into_parts(); + let ts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ts.value(0), 1); + assert_eq!(ts.value(1), -2); + + // Millisecond-encoded rows aligned to a microsecond schema: widening + // is lossless. + let mut rows = Rows { + schema: vec![rows.schema[0].clone(), rows.schema[1].clone()], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(123)), + }, + Value { + value_data: Some(ValueData::F64Value(4.0)), + }, + ], + }], + }; + rows.schema[0].datatype = ColumnDataType::TimestampMillisecond as i32; + let target = ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("greptime_value", DataType::Float64, true), + ]); + + let (batch, _) = rows_to_aligned_record_batch(&rows, &target) + .unwrap() + .into_parts(); + let ts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ts.value(0), 123_000); + } + #[test] fn test_rows_to_aligned_record_batch_renames_and_reorders() { let rows = Rows { diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 9182f64ad00..ff82f8f933b 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -18,20 +18,26 @@ use std::cmp::Ordering; use std::collections::HashMap; use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; +use std::sync::Arc; use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{Label, Query, Sample, TimeSeries, WriteRequest}; use api::v1::RowInsertRequests; use arrow::array::{ - Array, AsArray, DictionaryArray, LargeStringArray, StringArray, StringViewArray, + Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, StringArray, StringViewArray, +}; +use arrow::compute::kernels::cast as casts; +use arrow::datatypes::{ + DataType, Float64Type, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, UInt32Type, }; -use arrow::datatypes::{Float64Type, TimestampMillisecondType, UInt32Type}; use common_grpc::precision::Precision; +use common_memory_manager::MemoryGuard; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::{RecordBatch, RecordBatches}; use common_telemetry::{tracing, warn}; use datafusion::dataframe::DataFrame; -use datafusion::prelude::{Expr, col, lit, regexp_match}; +use datafusion::prelude::{Expr, ident, lit, regexp_match}; use datafusion_common::ScalarValue; use datafusion_expr::LogicalPlan; use snafu::{OptionExt, ResultExt, ensure}; @@ -39,6 +45,8 @@ use snap::raw::{Decoder, Encoder}; use crate::error::{self, Result}; use crate::prom_remote_write::REMOTE_WRITE_V1_VERSION; +use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; use crate::row_writer::{self, MultiTableData}; pub const METRIC_NAME_LABEL: &str = "__name__"; @@ -145,8 +153,9 @@ pub fn query_to_plan( let mut conditions = Vec::with_capacity(label_matches.len() + 1); conditions - .push(col(timestamp_column_name).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); - conditions.push(col(timestamp_column_name).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); + .push(ident(timestamp_column_name).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); + conditions + .push(ident(timestamp_column_name).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); for m in label_matches { let name = &m.name; @@ -165,18 +174,18 @@ pub fn query_to_plan( match m_type { MatcherType::Eq => { - conditions.push(col(name).eq(lit(value))); + conditions.push(ident(name).eq(lit(value))); } MatcherType::Neq => { - conditions.push(col(name).not_eq(lit(value))); + conditions.push(ident(name).not_eq(lit(value))); } // Case sensitive regexp match MatcherType::Re => { - conditions.push(regexp_match(col(name), lit(value), None).is_not_null()); + conditions.push(regexp_match(ident(name), lit(value), None).is_not_null()); } // Case sensitive regexp not match MatcherType::Nre => { - conditions.push(regexp_match(col(name), lit(value), None).is_null()); + conditions.push(regexp_match(ident(name), lit(value), None).is_null()); } } } @@ -398,6 +407,50 @@ fn append_recordbatch_to_timeseries( .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { msg: format!("missing timestamp column '{timestamp_column_name}' in query result"), })?; + // The Prometheus remote read wire format carries millisecond timestamps, + // while the table's time index can use any time unit (e.g. a metric + // physical table created with TIMESTAMP(6)). Narrowing floors towards + // negative infinity, consistent with `Timestamp::convert_to` on the + // ingestion path; arrow's cast would truncate towards zero, so a + // pre-epoch sample (e.g. -1001us) would round to -1ms instead of -2ms + // and disagree with what ingesting the same instant into a millisecond + // table would have stored. + let ts_column: ArrayRef = match ts_column.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, _) => ts_column.clone(), + // Second -> millisecond is a widening (exact) conversion. + DataType::Timestamp(TimeUnit::Second, _) => casts::cast( + ts_column, + &DataType::Timestamp(TimeUnit::Millisecond, None), + ) + .map_err(|e| { + error::InvalidPromRemoteReadQueryResultSnafu { + msg: format!( + "failed to cast timestamp column '{timestamp_column_name}' of datatype {:?} to millisecond: {e}", + ts_column.data_type() + ), + } + .build() + })?, + DataType::Timestamp(TimeUnit::Microsecond, _) => Arc::new( + ts_column + .as_primitive::() + .unary::<_, TimestampMillisecondType>(|v| v.div_euclid(1_000)), + ), + DataType::Timestamp(TimeUnit::Nanosecond, _) => Arc::new( + ts_column + .as_primitive::() + .unary::<_, TimestampMillisecondType>(|v| v.div_euclid(1_000_000)), + ), + _ => { + return error::InvalidPromRemoteReadQueryResultSnafu { + msg: format!( + "Expect timestamp column of datatype Timestamp(Millisecond), actual {:?}", + ts_column.data_type() + ), + } + .fail(); + } + }; let ts_column = ts_column .as_primitive_opt::() .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { @@ -553,8 +606,147 @@ pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertR Ok(multi_table_data.into_row_insert_requests()) } +/// Hard cap on the decompressed size of a handler-decompressed request body +/// (Prometheus remote read/write, Loki push). +/// +/// The HTTP body limit only bounds the *encoded* bytes on the wire. Without +/// this cap a tiny compressed payload (e.g. a five-byte raw-Snappy varint or a +/// few KiB of zstd RLE blocks) could declare/expand to multiple GiB of heap +/// before any protobuf validation runs, so the decompressed size must be +/// bounded on its own. +pub const MAX_DECOMPRESSED_REQUEST_SIZE: usize = 512 * 1024 * 1024; + +/// Size of the chunks a zstd body is decompressed and charged in. +const ZSTD_DECOMPRESS_CHUNK: usize = 256 * 1024; + +/// Decompressed bytes together with the memory permits charged for them. +/// The permits are released when the buffer is dropped, so the aggregate +/// memory accounting covers the whole lifetime of the decompressed data. +pub(crate) struct ChargedBuffer { + pub data: Vec, + guards: Vec>, +} + +impl std::ops::Deref for ChargedBuffer { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + &self.data + } +} + +impl std::fmt::Debug for ChargedBuffer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChargedBuffer") + .field("data_len", &self.data.len()) + .field("guards", &self.guards.len()) + .finish() + } +} + +impl ChargedBuffer { + /// Creates a buffer with the permits charged for it. + pub(crate) fn new(data: Vec, guards: Vec>) -> Self { + Self { data, guards } + } + + /// Splits the buffer into the raw bytes and the memory permits charged for + /// them. Hold the permits as long as the bytes are kept. + pub(crate) fn into_parts(self) -> (Vec, Vec>) { + (self.data, self.guards) + } +} + +/// Decompresses a raw Snappy body after enforcing a hard cap on the decoded +/// size and charging the aggregate request-memory limiter. +/// +/// The raw Snappy format prefixes the stream with a varint of the decoded +/// length, so the output size is known (and validated) before the output +/// buffer is allocated by `decompress_vec`. +pub(crate) async fn snappy_decompress_limited( + buf: &[u8], + max_decompressed: usize, + limiter: &ServerMemoryLimiter, +) -> Result { + let decoded_len = + snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?; + ensure!( + decoded_len <= max_decompressed, + error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: max_decompressed as u64, + } + ); + // Reserve the declared size before the allocation happens so concurrent + // decompressions cannot exceed the aggregate quota. + let guard = limiter.acquire(decoded_len as u64).await?; + let mut decoder = Decoder::new(); + let data = decoder + .decompress_vec(buf) + .context(error::DecompressSnappyPromRemoteRequestSnafu)?; + Ok(ChargedBuffer { + data, + guards: vec![guard], + }) +} + +/// Decompresses a zstd body with a hard cap on the decoded size, charging the +/// aggregate request-memory limiter incrementally while the output is +/// materialized (zstd streams don't declare their decoded size up front). +pub(crate) async fn zstd_decompress_limited( + buf: &[u8], + max_decompressed: usize, + limiter: &ServerMemoryLimiter, +) -> Result { + use std::io::Read; + + let decoder = zstd::stream::read::Decoder::new(buf) + .context(error::DecompressZstdPromRemoteRequestSnafu)?; + // Allow reading one byte past the limit so oversized streams are detected + // instead of being silently truncated. + let mut limited = decoder.take(max_decompressed as u64 + 1); + let mut data = Vec::new(); + let mut guards = Vec::new(); + let mut chunk = vec![0u8; ZSTD_DECOMPRESS_CHUNK]; + loop { + let n = limited + .read(&mut chunk) + .context(error::DecompressZstdPromRemoteRequestSnafu)?; + if n == 0 { + break; + } + ensure!( + data.len() + n <= max_decompressed, + error::DecompressedBodyTooLargeSnafu { + size: (data.len() + n) as u64, + limit: max_decompressed as u64, + } + ); + // Charge each chunk before it is appended so the aggregate quota + // bounds the memory while it is being materialized. + guards.push(limiter.acquire(n as u64).await?); + data.extend_from_slice(&chunk[..n]); + } + Ok(ChargedBuffer { data, guards }) +} + +/// Decompresses a raw Snappy body with the hard decoded-size cap enforced, +/// without charging the aggregate request-memory limiter. +/// +/// Used for response bodies (remote read replies) and tests; request bodies +/// must use [`snappy_decompress_limited`] instead so their decoded size is +/// also charged to the aggregate quota. #[inline] pub fn snappy_decompress(buf: &[u8]) -> Result> { + let decoded_len = + snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?; + ensure!( + decoded_len <= MAX_DECOMPRESSED_REQUEST_SIZE, + error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: MAX_DECOMPRESSED_REQUEST_SIZE as u64, + } + ); let mut decoder = Decoder::new(); decoder .decompress_vec(buf) @@ -569,11 +761,6 @@ pub fn snappy_compress(buf: &[u8]) -> Result> { .context(error::CompressPromRemoteRequestSnafu) } -#[inline] -pub fn zstd_decompress(buf: &[u8]) -> Result> { - zstd::stream::decode_all(buf).context(error::DecompressZstdPromRemoteRequestSnafu) -} - /// Mock timeseries for test, it is both used in servers and frontend crate /// So we present it here pub fn mock_timeseries() -> Vec { @@ -713,6 +900,7 @@ mod tests { DictionaryArray, Float64Array, StringArray, TimestampMillisecondArray, UInt32Array, }; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, UInt32Type}; + use common_memory_manager::OnExhaustedPolicy; use common_recordbatch::DfRecordBatch; use datafusion::prelude::SessionContext; use datatypes::data_type::ConcreteDataType; @@ -841,6 +1029,8 @@ mod tests { ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true), ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true), ColumnSchema::new("job", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("service.name", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("Region", ConcreteDataType::string_datatype(), true), ])); let recordbatch = RecordBatch::new( schema, @@ -849,6 +1039,8 @@ mod tests { Arc::new(Float64Vector::from_vec(vec![3.0])) as _, Arc::new(StringVector::from(vec!["host1"])) as _, Arc::new(StringVector::from(vec!["job"])) as _, + Arc::new(StringVector::from(vec!["api"])) as _, + Arc::new(StringVector::from(vec!["us"])) as _, ], ) .unwrap(); @@ -887,6 +1079,16 @@ mod tests { value: "localhost".to_string(), r#type: NEQ_TYPE, }, + LabelMatcher { + name: "service.name".to_string(), + value: "api".to_string(), + r#type: EQ_TYPE, + }, + LabelMatcher { + name: "Region".to_string(), + value: "us".to_string(), + r#type: EQ_TYPE, + }, ], ..Default::default() }; @@ -897,7 +1099,7 @@ mod tests { let ts_col = greptime_timestamp(); let expected = format!( - "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None) AND regexp_match(?table?.job, Utf8(\"*prom*\")) IS NOT NULL AND ?table?.instance != Utf8(\"localhost\")\n TableScan: ?table?", + "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None) AND regexp_match(?table?.job, Utf8(\"*prom*\")) IS NOT NULL AND ?table?.instance != Utf8(\"localhost\") AND ?table?.service.name = Utf8(\"api\") AND ?table?.Region = Utf8(\"us\")\n TableScan: ?table?", ts_col, ts_col ); assert_eq!(expected, display_string); @@ -1457,4 +1659,126 @@ mod tests { timeseries[1].samples ); } + + #[tokio::test] + async fn test_snappy_varint_bomb_is_rejected_before_allocation() { + // A five-byte raw Snappy header declaring a ~4 GiB decoded length. + // Regression test for the compression-bomb finding: the declared size + // must be rejected by the cap check before any buffer is allocated. + let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f]; + let limiter = ServerMemoryLimiter::default(); + + let err = snappy_decompress_limited(&bomb, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_snappy_decompress_limited_roundtrip() { + let raw = vec![b'a'; 1024 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + let limiter = ServerMemoryLimiter::default(); + + let decoded = + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + assert_eq!(&decoded[..], &raw[..]); + } + + #[tokio::test] + async fn test_snappy_charges_and_releases_quota() { + let raw = vec![b'a'; 64 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + let limiter = ServerMemoryLimiter::new(128 * 1024, OnExhaustedPolicy::Fail); + + let decoded = + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + // Rounding to the 1 KiB permit granularity. + assert!(limiter.used_bytes() >= 64 * 1024); + assert!(limiter.used_bytes() <= 65 * 1024); + drop(decoded); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_snappy_rejects_when_quota_exhausted() { + let raw = vec![b'a'; 64 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + // Quota smaller than the decoded size. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + + let err = snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::MemoryLimitExceeded { .. } + )); + assert_eq!(0, limiter.used_bytes(), "failed charge must be released"); + } + + #[tokio::test] + async fn test_zstd_at_and_over_the_cap() { + let limiter = ServerMemoryLimiter::default(); + let small_cap = 8 * 1024; + + // Exactly at the cap: accepted. + let raw = vec![0u8; small_cap]; + let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); + let decoded = zstd_decompress_limited(&compressed, small_cap, &limiter) + .await + .unwrap(); + assert_eq!(decoded.len(), small_cap); + + // One byte past the cap: rejected. + let raw = vec![0u8; small_cap + 1]; + let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); + let err = zstd_decompress_limited(&compressed, small_cap, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + } + + #[tokio::test] + async fn test_zstd_charges_incrementally() { + // 64 KiB wire that expands to 2 MiB must be charged chunk by chunk + // against a quota that only fits the expanded size. + let raw = vec![0u8; 2 * 1024 * 1024]; + let compressed = zstd::stream::encode_all(&raw[..], 19).unwrap(); + assert!(compressed.len() < 16 * 1024); + let limiter = ServerMemoryLimiter::new(4 * 1024 * 1024, OnExhaustedPolicy::Fail); + + let decoded = zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + assert_eq!(decoded.len(), raw.len()); + assert!(limiter.used_bytes() >= raw.len() as u64); + drop(decoded); + assert_eq!(0, limiter.used_bytes()); + } + + #[test] + fn test_snappy_response_decompress_is_capped() { + // The response-side (uncharged) helper must enforce the cap too. + let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f]; + let err = snappy_decompress(&bomb).unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + } } diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index e74088a2a5d..12a56076d6f 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -81,8 +81,6 @@ pub struct MetricsIngestOutcome { #[async_trait] pub trait InfluxdbLineProtocolHandler { - /// A successful request will not return a response. - /// Only on error will the socket return a line of data. async fn exec(&self, request: InfluxdbRequest, ctx: QueryContextRef) -> Result; } @@ -91,12 +89,10 @@ pub trait OpentsdbProtocolHandler { /// Checks all points in one external request before per-point debug execution. async fn preflight(&self, data_points: &[DataPoint], ctx: QueryContextRef) -> Result<()>; - /// A successful request will not return a response. - /// Only on error will the socket return a line of data. async fn exec(&self, data_points: Vec, ctx: QueryContextRef) -> Result; - /// Executes an ordinary HTTP put with optional batching. Debug and socket - /// callers retain [`Self::exec`]; the frontend clears HTTP batching selection + /// Executes an ordinary HTTP put with optional batching. Debug callers + /// retain [`Self::exec`]; the frontend clears HTTP batching selection /// there so diagnostic requests preserve direct, per-point error attribution. async fn exec_batch(&self, data_points: Vec, ctx: QueryContextRef) -> Result { self.exec(data_points, ctx).await diff --git a/src/servers/src/query_handler/sql.rs b/src/servers/src/query_handler/sql.rs index 9bf99a44441..1b65ff002db 100644 --- a/src/servers/src/query_handler/sql.rs +++ b/src/servers/src/query_handler/sql.rs @@ -26,8 +26,30 @@ use crate::error::Result; pub type ServerSqlQueryHandlerRef = Arc; +/// Maximum decoded SQL size accepted by the logical-table batch operation. +pub const MAX_LOGICAL_TABLE_DDL_BYTES: usize = 1024 * 1024; +/// Maximum number of CREATE statements in one logical-table batch. +pub const MAX_LOGICAL_TABLE_DDL_STATEMENTS: usize = 128; + #[async_trait] pub trait SqlQueryHandler { + /// Whether this handler implements the bounded logical-table DDL operation. + fn supports_metric_batch_ddl(&self) -> bool { + false + } + + /// Validates and authorizes the whole batch before submitting one procedure. + async fn create_logical_tables( + &self, + _query: &str, + _query_ctx: QueryContextRef, + ) -> Result> { + crate::error::UnsupportedOperationSnafu { + operation: "metric batch DDL", + } + .fail() + } + async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec>; /// Executes an HTTP analyze-stream query. diff --git a/src/servers/src/request_memory_limiter.rs b/src/servers/src/request_memory_limiter.rs index 8c3d106d9c0..ac072f51b2e 100644 --- a/src/servers/src/request_memory_limiter.rs +++ b/src/servers/src/request_memory_limiter.rs @@ -73,4 +73,9 @@ impl ServerMemoryLimiter { pub fn limit_bytes(&self) -> u64 { self.manager.limit_bytes() } + + /// Returns currently reserved (in-flight) bytes. + pub fn used_bytes(&self) -> u64 { + self.manager.used_bytes() + } } diff --git a/src/servers/tests/http/mod.rs b/src/servers/tests/http/mod.rs index cca2dfd7873..ffd76b3753b 100644 --- a/src/servers/tests/http/mod.rs +++ b/src/servers/tests/http/mod.rs @@ -16,4 +16,5 @@ mod authorize; mod http_handler_test; mod influxdb_test; mod opentsdb_test; +mod prom_count_values_test; mod prom_store_test; diff --git a/src/servers/tests/http/prom_count_values_test.rs b/src/servers/tests/http/prom_count_values_test.rs new file mode 100644 index 00000000000..ab17ef8791f --- /dev/null +++ b/src/servers/tests/http/prom_count_values_test.rs @@ -0,0 +1,192 @@ +// 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. + +//! Regression tests for the Prometheus HTTP API conversion of `count_values` results. +//! +//! `count_values` turns a sample value into a label, so the generated column must be a +//! *string* column holding Prometheus' textual form of the value: the conversion infers +//! labels from string columns and the sample value from the first numeric column. A numeric +//! label column is reported as the sample value of the series instead of a label. + +use std::sync::Arc; + +use catalog::RegisterTableRequest; +use catalog::memory::MemoryCatalogManager; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_query::prelude::greptime_value; +use common_recordbatch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::{Float64Vector, StringVector, TimestampMillisecondVector}; +use promql_parser::parser::value::ValueType; +use query::options::QueryOptions; +use query::parser::{PromQuery, QueryLanguageParser}; +use query::query_engine::QueryEngineFactory; +use servers::http::prometheus::{ + PromData, PromQueryResult, PrometheusJsonResponse, PrometheusResponse, +}; +use session::context::QueryContext; +use table::Table; +use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; +use table::test_util::MemTable; + +/// Catalog with a single metric `cv_metric` holding one series per value at `ts=1000ms`. +fn metric_catalog_manager(values: &[f64]) -> Arc { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(3_001) + .name("cv_metric") + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(StringVector::from( + (0..values.len()) + .map(|index| format!("k{index}")) + .collect::>(), + )) as _, + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000; + values.len() + ])) as _, + Arc::new(Float64Vector::from_vec(values.to_vec())) as _, + ], + ) + .unwrap(); + let backing = MemTable::new_with_catalog( + "cv_metric", + batch, + 3_001, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "cv_metric".to_string(), + table_id: 3_001, + table, + }) + .is_ok() + ); + + catalog_list +} + +#[tokio::test] +async fn count_values_generated_label_is_reported_as_a_prometheus_label() { + // Two samples carry the value 5.0, one carries 0.5 and one carries 200.0. The query + // evaluates at `ts=1s`, the only timestamp holding samples. + let catalog_list = metric_catalog_manager(&[5.0, 5.0, 0.5, 200.0]); + let query_engine = QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + + let query_ctx = QueryContext::arc(); + let prom_query = PromQuery { + query: r#"count_values("v", cv_metric) + 1"#.to_string(), + start: "1".to_string(), + end: "1".to_string(), + step: "5s".to_string(), + lookback: "5m".to_string(), + alias: None, + }; + let statement = QueryLanguageParser::parse_promql(&prom_query, &query_ctx).unwrap(); + let plan = query_engine + .planner() + .plan(&statement, query_ctx.clone()) + .await + .unwrap(); + let output = query_engine.execute(plan, query_ctx).await.unwrap(); + + let response = + PrometheusJsonResponse::from_query_result(Ok(output), None, ValueType::Vector, None).await; + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = response.data + else { + panic!("expected a vector response"); + }; + + // `v` is the generated label, so it must be reported as a label holding Prometheus' + // textual form of the sample value (`5`, not `5.0`), and the series carries the + // aggregation result (`count_values("v", cv_metric) + 1`) as its sample value. + let mut actual = series + .into_iter() + .map(|series| { + assert_eq!( + series.metric.keys().collect::>(), + vec!["v"], + "`v` must be the only label of the series" + ); + let label = series.metric["v"].clone(); + let (timestamp, value) = series + .value + .expect("every series of a vector result carries a sample"); + (label, timestamp, value) + }) + .collect::>(); + actual.sort_by(|left, right| left.0.cmp(&right.0)); + + assert_eq!( + actual, + vec![ + ("0.5".to_string(), 1.0, "2.0".to_string()), + ("200".to_string(), 1.0, "2.0".to_string()), + ("5".to_string(), 1.0, "3.0".to_string()), + ] + ); +} diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index 7c07e3b0344..c7a288ccee0 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -36,11 +36,12 @@ use datafusion_expr::LogicalPlan; use prost::Message; use query::parser::PromQuery; use query::query_engine::DescribeResult; +use servers::batcher::BatchingProtocol; use servers::error::{self, Result}; use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF}; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::http::test_helpers::{TestClient, TestResponse}; -use servers::http::{BatchingProtocol, HttpOptions, HttpServerBuilder}; +use servers::http::{HttpOptions, HttpServerBuilder}; use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store; @@ -208,34 +209,10 @@ fn make_test_app_with_write_capture( make_test_app_with_write_failure(read_tx, write_tx, None) } -fn make_test_app_with_native_histogram_write_capture( - read_tx: mpsc::Sender<(String, Vec)>, - write_tx: mpsc::Sender, -) -> Router { - make_test_app_with_write_failure_inner(read_tx, write_tx, None, true) -} - fn make_test_app_with_write_failure( read_tx: mpsc::Sender<(String, Vec)>, write_tx: mpsc::Sender, fail_write_call: Option, -) -> Router { - make_test_app_with_write_failure_inner(read_tx, write_tx, fail_write_call, false) -} - -fn make_test_app_with_native_histogram_write_failure( - read_tx: mpsc::Sender<(String, Vec)>, - write_tx: mpsc::Sender, - fail_write_call: Option, -) -> Router { - make_test_app_with_write_failure_inner(read_tx, write_tx, fail_write_call, true) -} - -fn make_test_app_with_write_failure_inner( - read_tx: mpsc::Sender<(String, Vec)>, - write_tx: mpsc::Sender, - fail_write_call: Option, - experimental_enable_prometheus_native_histogram: bool, ) -> Router { let http_opts = HttpOptions { addr: format!("127.0.0.1:{}", ports::get_port()), @@ -261,14 +238,7 @@ fn make_test_app_with_write_failure_inner( BatchingProtocol::HttpSql, ]) .with_sql_handler(instance.clone()) - .with_prom_handler( - instance, - None, - true, - PromValidationMode::Unchecked, - experimental_enable_prometheus_native_histogram, - None, - ) + .with_prom_handler(instance, None, true, PromValidationMode::Unchecked, None) .build(); server.build(server.make_app()).unwrap() } @@ -417,7 +387,7 @@ async fn test_prometheus_remote_write_v2_histogram_write_error_has_partial_writt let (read_tx, _read_rx) = mpsc::channel(100); let (write_tx, mut write_rx) = mpsc::channel(100); - let app = make_test_app_with_native_histogram_write_failure(read_tx, write_tx, Some(2)); + let app = make_test_app_with_write_failure(read_tx, write_tx, Some(2)); let client = TestClient::new(app).await; let mut write_request = remote_write_v2::request_with_labels_and_samples( @@ -670,7 +640,7 @@ async fn test_prometheus_remote_write_v2_writes_histogram_only_series() { let (read_tx, _read_rx) = mpsc::channel(100); let (write_tx, mut write_rx) = mpsc::channel(100); - let app = make_test_app_with_native_histogram_write_capture(read_tx, write_tx); + let app = make_test_app_with_write_capture(read_tx, write_tx); let client = TestClient::new(app).await; let write_request = remote_write_v2::request_with_labels_and_histograms( @@ -707,36 +677,6 @@ async fn test_prometheus_remote_write_v2_writes_histogram_only_series() { assert!(write_rx.try_recv().is_err()); } -#[tokio::test] -async fn test_prometheus_remote_write_v2_rejects_native_histogram_when_disabled() { - common_telemetry::init_default_ut_logging(); - let (read_tx, _read_rx) = mpsc::channel(100); - let (write_tx, mut write_rx) = mpsc::channel(100); - - let app = make_test_app_with_write_capture(read_tx, write_tx); - let client = TestClient::new(app).await; - - let write_request = remote_write_v2::request_with_labels_and_histograms( - vec![( - prom_store::METRIC_NAME_LABEL, - "http_request_duration_seconds", - )], - vec![remote_write_v2::histogram(1000)], - ); - - let result = post_remote_write_v2(&client, &write_request).await; - - assert_eq!(result.status(), 400); - assert_remote_write_v2_written_headers_with_histograms(&result.headers(), "0", "0"); - assert!( - result - .text() - .await - .contains("prom_store.experimental_enable_prometheus_native_histogram") - ); - assert!(write_rx.try_recv().is_err()); -} - #[tokio::test] async fn test_prometheus_remote_write_rejects_unsupported_proto() { common_telemetry::init_default_ut_logging(); @@ -872,7 +812,6 @@ async fn test_prom_batching_depends_on_protocol_and_metric_engine() { None, with_metric_engine, PromValidationMode::Unchecked, - false, None, ) .build(); diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 2e031c0d0db..6b2995283f8 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -587,6 +587,92 @@ async fn start_test_server(server_tls: TlsOption) -> Result { Ok(server_addr.port()) } +const EMPTY_QUERIES: &[&str] = &[ + "", + " \t\r\n", + ";", + ";;;", + "-- ping", + "-- ping\n", + "/* ping */", + "/* outer /* inner */ comment */", + "; -- ping\r\n /* comment */ ;", +]; + +#[tokio::test] +async fn test_simple_query_empty_statements() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for query in EMPTY_QUERIES { + let messages = client + .simple_query(query) + .await + .unwrap_or_else(|err| panic!("query {query:?} failed: {err}")); + // tokio-postgres exposes EmptyQueryResponse as CommandComplete(0). + assert!( + matches!( + messages.as_slice(), + [SimpleQueryMessage::CommandComplete(0)] + ), + "expected an empty query response for {query:?}, got {messages:?}" + ); + } + + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert_eq!(unwrap_results(&messages), vec!["1"]); + Ok(()) +} + +#[tokio::test] +async fn test_extended_query_empty_statements() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for query in EMPTY_QUERIES { + let statement = client + .prepare(query) + .await + .unwrap_or_else(|err| panic!("prepare {query:?} failed: {err}")); + assert_eq!(client.execute(&statement, &[]).await.unwrap(), 0); + assert!(client.query(&statement, &[]).await.unwrap().is_empty()); + } + + let row = client + .query_one("SELECT $1::BIGINT", &[&42i64]) + .await + .unwrap(); + assert_eq!(row.get::<_, i64>(0), 42); + Ok(()) +} + +#[tokio::test] +async fn test_simple_query_with_comments() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for (query, expected) in [ + ("-- ping\nSELECT 1", vec!["1"]), + ("/* comment */ SELECT 1; -- trailing", vec!["1"]), + ("SELECT '-- ping'", vec!["-- ping"]), + ("SELECT '/* comment */'", vec!["/* comment */"]), + ("SELECT 1; /* between */ SELECT 2;", vec!["1", "2"]), + ] { + let messages = client.simple_query(query).await.unwrap(); + assert_eq!(unwrap_results(&messages), expected, "query: {query}"); + } + + assert!( + client + .simple_query("/* comment */ SELECT missing_column FROM numbers") + .await + .is_err() + ); + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert_eq!(unwrap_results(&messages), vec!["1"]); + Ok(()) +} + async fn do_simple_query(server_tls: TlsOption, client_tls: bool) -> Result<()> { let server_port = start_test_server(server_tls).await?; diff --git a/src/servers/tests/prom_remote_write_v2_test.rs b/src/servers/tests/prom_remote_write_v2_test.rs index 75fee046f7c..7855f92b814 100644 --- a/src/servers/tests/prom_remote_write_v2_test.rs +++ b/src/servers/tests/prom_remote_write_v2_test.rs @@ -77,7 +77,7 @@ fn test_decode_remote_write_v2_native_histogram_dump() { assert_eq!(histogram.timestamp, 1782358160412); let (sample_inserts, histogram_inserts, sample_count, histogram_count) = - remote_write_v2::decode_write_requests(false, Bytes::from_static(BODY), true).unwrap(); + remote_write_v2::decode_write_requests(false, Bytes::from_static(BODY)).unwrap(); assert!(sample_inserts.is_empty()); assert_eq!(sample_count, 0); assert_eq!(histogram_count, 1); diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 889b911c3cc..93c46ff7ff7 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -4,9 +4,6 @@ version.workspace = true edition.workspace = true license.workspace = true -[features] -testing = [] - [lints] workspace = true @@ -29,3 +26,12 @@ derive_more.workspace = true snafu.workspace = true sql.workspace = true uuid.workspace = true + +[dev-dependencies] +prost.workspace = true + +[package.metadata.cargo-shear] +# `common-error` is required by `common-macro`'s `stack_trace_debug`: the +# generated code references `::common_error` paths, so it must stay a direct +# dependency. +ignored = ["common-error"] diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 284b0bdbfa8..4a75713b97b 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -15,6 +15,7 @@ use std::collections::HashMap; use std::fmt::{Display, Formatter}; use std::net::SocketAddr; +use std::str::FromStr; use std::sync::{Arc, RwLock}; use std::time::Duration; @@ -34,12 +35,14 @@ use derive_builder::Builder; use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect}; pub use crate::hints::{ - LIVE_ANALYZE_METRICS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, - SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + LIVE_ANALYZE_METRICS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY, + REMOTE_QUERY_ID_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, }; use crate::protocol_ctx::ProtocolCtx; use crate::query_id::QueryId; -use crate::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle}; +use crate::session_config::{ + InvalidConfigValueSnafu, PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle, +}; use crate::{MutableInner, ReadPreference}; pub type QueryContextRef = Arc; @@ -64,7 +67,6 @@ pub struct QueryContext { current_catalog: String, /// mapping of RegionId to SequenceNumber, for snapshot read, meaning that the read should only /// container data that was committed before(and include) the given sequence number - /// this field will only be filled if extensions contains a pair of "snapshot_read" and "true" snapshot_seqs: Arc>>, /// Mappings of the RegionId to the minimal sequence of SST file to scan. sst_min_sequences: Arc>>, @@ -84,6 +86,9 @@ pub struct QueryContext { /// Local-only write batching selection; never transported in protobuf extensions. #[builder(default)] batching_enabled: bool, + /// Local-only opt-in for logical metric writes, independent of ordinary-table batching. + #[builder(default)] + logical_batching_enabled: bool, /// Track which protocol the query comes from. #[builder(default)] channel: Channel, @@ -180,25 +185,41 @@ impl QueryContextBuilder { } } -impl From<&RegionRequestHeader> for QueryContext { - fn from(value: &RegionRequestHeader) -> Self { +impl TryFrom<&RegionRequestHeader> for QueryContext { + type Error = crate::session_config::Error; + + fn try_from(value: &RegionRequestHeader) -> Result { if let Some(ctx) = &value.query_context { - ctx.clone().into() + ctx.clone().try_into() } else { - QueryContextBuilder::default() + Ok(QueryContextBuilder::default() .set_extension( REMOTE_QUERY_ID_EXTENSION_KEY.to_string(), generate_remote_query_id(), ) - .build() + .build()) } } } -impl From for QueryContext { - fn from(ctx: api::v1::QueryContext) -> Self { +impl TryFrom for QueryContext { + type Error = crate::session_config::Error; + + fn try_from(mut ctx: api::v1::QueryContext) -> Result { + let read_preference = ctx.extensions.remove(READ_PREFERENCE_EXTENSION_KEY); + let read_preference = match read_preference { + Some(value) => ReadPreference::from_str(&value).map_err(|_| { + InvalidConfigValueSnafu { + name: "read_preference", + value, + hint: "Expected a supported read preference", + } + .build() + })?, + None => ReadPreference::default(), + }; let sequences = ctx.snapshot_seqs.as_ref(); - QueryContextBuilder::default() + Ok(QueryContextBuilder::default() .current_catalog(ctx.current_catalog) .current_schema(ctx.current_schema) .timezone(parse_timezone(Some(&ctx.timezone))) @@ -215,7 +236,8 @@ impl From for QueryContext { .unwrap_or_default(), ))) .explain_options(ctx.explain) - .build() + .read_preference(read_preference) + .build()) } } @@ -234,6 +256,17 @@ impl From for api::v1::QueryContext { ) -> Self { let explain = mutable_query_context_data.read().unwrap().explain_options; let mutable_inner = mutable_inner.read().unwrap(); + let mut extensions = extensions; + extensions.remove(READ_PREFERENCE_EXTENSION_KEY); + if mutable_inner.read_preference != ReadPreference::Leader { + extensions.insert( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + mutable_inner + .read_preference + .to_string() + .to_ascii_lowercase(), + ); + } api::v1::QueryContext { current_catalog, current_schema: mutable_inner.schema.clone(), @@ -507,7 +540,17 @@ impl QueryContext { &self.configuration_parameter } - /// Whether the local HTTP entry point selected write batching. + /// Whether the local HTTP entry point selected logical-table batching. + pub fn logical_batching_enabled(&self) -> bool { + self.logical_batching_enabled + } + + /// Sets local logical-table batching selection without adding a wire-visible extension. + pub fn set_logical_batching_enabled(&mut self, enabled: bool) { + self.logical_batching_enabled = enabled; + } + + /// Whether the local protocol entry point selected ordinary-table batching. pub fn batching_enabled(&self) -> bool { self.batching_enabled } @@ -684,6 +727,7 @@ impl QueryContextBuilder { channel, batching_enabled: self.batching_enabled.unwrap_or_default(), admitted_write: None, + logical_batching_enabled: self.logical_batching_enabled.unwrap_or_default(), process_id: self.process_id.unwrap_or_default(), conn_info: self.conn_info.unwrap_or_default(), protocol_ctx: self.protocol_ctx.unwrap_or_default(), @@ -908,7 +952,7 @@ mod test { 0 ); let wire: api::v1::QueryContext = admitted.into(); - let restored = QueryContext::from(wire); + let restored = QueryContext::try_from(wire).unwrap(); assert_eq!(restored.write_rows_to_admit("greptime", "public", 100), 100); } @@ -921,7 +965,7 @@ mod test { .extensions .contains_key(crate::hints::INSERT_SKIP_WAL_HINT) ); - let restored: QueryContext = api_context.into(); + let restored: QueryContext = api_context.try_into().unwrap(); assert!(!restored.skip_wal()); } @@ -940,7 +984,7 @@ mod test { explain: None, }; - let session_ctx: QueryContext = api_ctx.clone().into(); + let session_ctx: QueryContext = api_ctx.clone().try_into().unwrap(); let roundtrip_api: api::v1::QueryContext = session_ctx.into(); assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog); @@ -975,7 +1019,7 @@ mod test { assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id); let proto: api::v1::QueryContext = (&ctx).into(); - let restored = QueryContext::from(proto); + let restored = QueryContext::try_from(proto).unwrap(); assert_eq!(restored.remote_query_id(), Some(query_id)); assert_eq!( restored.remote_query_id_value().unwrap().to_string(), @@ -983,6 +1027,117 @@ mod test { ); } + #[test] + fn test_query_context_read_preference_wire_values() { + let wire = |value: Option<&str>| api::v1::QueryContext { + current_catalog: DEFAULT_CATALOG_NAME.to_string(), + current_schema: DEFAULT_SCHEMA_NAME.to_string(), + timezone: "UTC".to_string(), + extensions: value + .map(|value| { + HashMap::from([(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string())]) + }) + .unwrap_or_default(), + channel: Channel::Internal as u32, + snapshot_seqs: None, + explain: None, + }; + + let header = RegionRequestHeader::default(); + let header_context = QueryContext::try_from(&header).unwrap(); + assert_eq!(header_context.read_preference(), ReadPreference::Leader); + + for value in [None, Some("leader"), Some("LEADER")] { + let ctx = QueryContext::try_from(wire(value)).unwrap(); + assert_eq!(ctx.read_preference(), ReadPreference::Leader); + assert!(!ctx.extensions().contains_key(READ_PREFERENCE_EXTENSION_KEY)); + } + for value in ["", "unknown", "follower", "follower_preferred", "random"] { + let error = QueryContext::try_from(wire(Some(value))).unwrap_err(); + assert!(matches!( + error, + crate::session_config::Error::InvalidConfigValue { + name, + value: actual, + .. + } if name == "read_preference" && actual == value + )); + } + } + + #[test] + fn test_legacy_query_context_wire_defaults_to_leader() { + use prost::Message; + + let legacy = api::v1::QueryContext::decode( + include_bytes!( + "../../../tests/compatibility/wire/read_preference/legacy_query_context.pb" + ) + .as_slice(), + ) + .unwrap(); + assert_eq!(legacy.current_catalog, "c1"); + assert_eq!(legacy.current_schema, "s1"); + assert_eq!(legacy.timezone, "UTC"); + assert_eq!(legacy.channel, Channel::Internal as u32); + assert_eq!( + legacy + .extensions + .get("flow.return_region_seq") + .map(String::as_str), + Some("true") + ); + assert!( + !legacy + .extensions + .contains_key(READ_PREFERENCE_EXTENSION_KEY) + ); + + let restored = QueryContext::try_from(legacy).unwrap(); + assert_eq!(restored.read_preference(), ReadPreference::Leader); + assert_eq!(restored.current_catalog(), "c1"); + assert_eq!(restored.current_schema(), "s1"); + assert_eq!(restored.timezone().to_string(), "UTC"); + assert_eq!(restored.channel(), Channel::Internal); + assert_eq!(restored.extension("flow.return_region_seq"), Some("true")); + + let outbound: api::v1::QueryContext = restored.into(); + let roundtrip = api::v1::QueryContext::decode(outbound.encode_to_vec().as_slice()).unwrap(); + assert_eq!(roundtrip.current_catalog, "c1"); + assert_eq!(roundtrip.current_schema, "s1"); + assert_eq!(roundtrip.timezone, "UTC"); + assert_eq!(roundtrip.channel, Channel::Internal as u32); + assert_eq!( + roundtrip + .extensions + .get("flow.return_region_seq") + .map(String::as_str), + Some("true") + ); + assert!( + !roundtrip + .extensions + .contains_key(READ_PREFERENCE_EXTENSION_KEY) + ); + } + + #[test] + fn test_outbound_read_preference_clears_stale_extension() { + let context = QueryContextBuilder::default() + .set_extension( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + "stale".to_string(), + ) + .set_extension("unrelated".to_string(), "preserved".to_string()) + .build(); + let wire: api::v1::QueryContext = context.into(); + assert!(!wire.extensions.contains_key(READ_PREFERENCE_EXTENSION_KEY)); + assert_eq!( + wire.extensions.get("unrelated").map(String::as_str), + Some("preserved") + ); + } + #[test] fn test_live_analyze_metrics_requires_matching_remote_query_id() { let mut ctx = QueryContext::arc().as_ref().clone(); @@ -1001,11 +1156,17 @@ mod test { fn test_batching_selection_is_local_only() { let mut ctx = QueryContextBuilder::default().build(); assert!(!ctx.batching_enabled()); + assert!(!ctx.logical_batching_enabled()); + ctx.set_logical_batching_enabled(true); + assert!(ctx.clone().logical_batching_enabled()); + assert!(ctx.fork().logical_batching_enabled()); ctx.set_batching_enabled(true); assert!(ctx.clone().batching_enabled()); assert!(ctx.fork().batching_enabled()); let wire: api::v1::QueryContext = ctx.into(); - assert!(!QueryContext::from(wire).batching_enabled()); + let restored = QueryContext::try_from(wire).unwrap(); + assert!(!restored.batching_enabled()); + assert!(!restored.logical_batching_enabled()); let ctx = QueryContextBuilder::default() .set_extension("batching_enabled".to_string(), "true".to_string()) .build(); diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index 971426f4cc0..3ed4fd8dc2e 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -27,11 +27,13 @@ pub const LIVE_ANALYZE_METRICS_EXTENSION_KEY: &str = "query.live_analyze_metrics 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] = [ +pub const READ_PREFERENCE_EXTENSION_KEY: &str = "query.read_preference"; +pub const RESERVED_EXTENSION_KEYS: [&str; 5] = [ REMOTE_QUERY_ID_EXTENSION_KEY, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, LIVE_ANALYZE_METRICS_EXTENSION_KEY, + READ_PREFERENCE_EXTENSION_KEY, ]; /// Deprecated, use `HINTS_KEY` instead. @@ -65,6 +67,7 @@ mod tests { assert!(is_reserved_extension_key( LIVE_ANALYZE_METRICS_EXTENSION_KEY )); + assert!(is_reserved_extension_key(READ_PREFERENCE_EXTENSION_KEY)); assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT)); } } diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index d9a4a386f91..5165def3e49 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -42,7 +42,6 @@ impl ProtocolCtx { /// If true, all scope attributes will be promoted to the final table schema. /// Along with the scope name, scope version and scope schema URL. /// - `with_metric_engine` -/// - `experimental_enable_exponential_histogram` /// - `is_legacy` /// If the user uses OTLP metrics ingestion before v0.16, it uses the old path. /// So we call this path 'legacy'. @@ -54,7 +53,6 @@ pub struct OtlpMetricCtx { pub resource_attrs: HashSet, pub promote_scope_attrs: bool, pub with_metric_engine: bool, - pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, /// Set from the server's `otlp.experimental_enable_resource_info`; off /// means the resource descriptor is not synthesized at all. diff --git a/src/sql/Cargo.toml b/src/sql/Cargo.toml index 1ab966fc7e3..de972bec30e 100644 --- a/src/sql/Cargo.toml +++ b/src/sql/Cargo.toml @@ -13,28 +13,23 @@ enterprise = [] [dependencies] api.workspace = true arrow-buffer.workspace = true +arrow-schema.workspace = true chrono.workspace = true common-base.workspace = true common-catalog.workspace = true -common-decimal.workspace = true common-error.workspace = true common-macro.workspace = true common-query.workspace = true common-sql.workspace = true -common-telemetry.workspace = true common-time.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true -datafusion-physical-expr.workspace = true datafusion-sql.workspace = true datatypes.workspace = true either.workspace = true -hex = "0.4" -humantime.workspace = true iso8601 = "0.6.1" itertools.workspace = true -jsonb.workspace = true lazy_static.workspace = true promql-parser.workspace = true regex.workspace = true @@ -43,7 +38,6 @@ serde_json.workspace = true snafu.workspace = true sqlparser.workspace = true sqlparser_derive = "0.1" -store-api.workspace = true table.workspace = true uuid.workspace = true diff --git a/src/sql/src/error.rs b/src/sql/src/error.rs index 1588842a365..f3699a70c39 100644 --- a/src/sql/src/error.rs +++ b/src/sql/src/error.rs @@ -294,13 +294,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Failed to set VECTOR index option"))] - SetVectorIndexOption { - source: datatypes::error::Error, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display( "Invalid partition number: {}, should be in range [2, 65536]", partition_num @@ -411,9 +404,7 @@ impl ErrorExt for Error { ConvertValue { .. } => StatusCode::Unsupported, PermissionDenied { .. } => StatusCode::PermissionDenied, - SetFulltextOption { .. } - | SetSkippingIndexOption { .. } - | SetVectorIndexOption { .. } => StatusCode::Unexpected, + SetFulltextOption { .. } | SetSkippingIndexOption { .. } => StatusCode::Unexpected, } } diff --git a/src/sql/src/lib.rs b/src/sql/src/lib.rs index b017f798f3f..e6bfe327902 100644 --- a/src/sql/src/lib.rs +++ b/src/sql/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(box_patterns)] - pub mod ast; pub mod dialect; pub mod error; diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index ca38ec0d65e..62c417d2537 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -26,7 +26,7 @@ use sqlparser::parser::IsOptional::Mandatory; use sqlparser::parser::{Parser, ParserError}; use sqlparser::tokenizer::{Token, TokenWithSpan}; -use crate::ast::ObjectNamePartExt; +use crate::ast::{ColumnOption, ObjectNamePartExt}; use crate::error::{self, InvalidColumnOptionSnafu, Result, SetFulltextOptionSnafu}; use crate::parser::ParserContext; use crate::parsers::create_parser::{INVERTED, parse_json2_type_and_options}; @@ -379,14 +379,67 @@ impl ParserContext<'_> { Ok(AlterTableOperation::AddConstraint(constraint)) } else { self.parser.prev_token(); - let add_columns = self - .parser - .parse_comma_separated(parse_add_columns) - .context(error::SyntaxSnafu)?; + let mut add_columns = vec![self.parse_add_column()?]; + while self.parser.consume_token(&Token::Comma) { + if matches!( + self.parser.peek_token().token, + Token::SemiColon | Token::EOF + ) || self.matches_keyword(Keyword::WITH) + { + break; + } + add_columns.push(self.parse_add_column()?); + } Ok(AlterTableOperation::AddColumns { add_columns }) } } + fn parse_add_column(&mut self) -> Result { + self.parser + .expect_keyword(Keyword::ADD) + .context(error::SyntaxSnafu)?; + let _ = self.parser.parse_keyword(Keyword::COLUMN); + let add_if_not_exists = + self.parser + .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); + let column = self.parse_column_def(true)?; + + // These attributes need ALTER-specific execution support. + ensure!( + column.extensions.fulltext_index_options.is_none() + && column.extensions.skipping_index_options.is_none() + && column.extensions.inverted_index_options.is_none() + && !column + .options() + .iter() + .any(|option| matches!(option.option, ColumnOption::DialectSpecific(_))), + InvalidColumnOptionSnafu { + name: column.name().to_string(), + msg: "ADD COLUMN does not support inline indexes or TIME INDEX", + } + ); + + let location = if self.parser.parse_keyword(Keyword::FIRST) { + Some(AddColumnLocation::First) + } else if matches!(self.parser.peek_token().token, Token::Word(word) if word.value.eq_ignore_ascii_case("AFTER")) + { + self.parser.next_token(); + let name = Self::canonicalize_identifier( + self.parser.parse_identifier().context(error::SyntaxSnafu)?, + ); + Some(AddColumnLocation::After { + column_name: name.value, + }) + } else { + None + }; + Ok(AddColumn { + column, + location, + add_if_not_exists, + }) + } + fn parse_alter_table_drop_default( &mut self, column_name: Ident, @@ -462,19 +515,19 @@ impl ParserContext<'_> { .context(error::SyntaxSnafu)?; self.parse_alter_table_drop_default(column_name) } else { - 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, - ) - }; + if let Some((_, json2_options)) = + parse_json2_type_and_options(&mut self.parser)? + { + return Ok(AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + }); + } + Ok(AlterTableOperation::ModifyColumnType { column_name, - target_type: data_type, - json2_options, + target_type: self.parser.parse_data_type().context(error::SyntaxSnafu)?, + json2_options: None, }) } } @@ -704,34 +757,6 @@ fn parse_string_options(parser: &mut Parser) -> std::result::Result<(String, Str Ok((name, value)) } -fn parse_add_columns(parser: &mut Parser) -> std::result::Result { - parser.expect_keyword(Keyword::ADD)?; - let _ = parser.parse_keyword(Keyword::COLUMN); - let add_if_not_exists = parser.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); - let mut column_def = parser.parse_column_def()?; - column_def.name = ParserContext::canonicalize_identifier(column_def.name); - let location = if parser.parse_keyword(Keyword::FIRST) { - Some(AddColumnLocation::First) - } else if let Token::Word(word) = parser.peek_token().token { - if word.value.eq_ignore_ascii_case("AFTER") { - let _ = parser.next_token(); - let name = ParserContext::canonicalize_identifier(parser.parse_identifier()?); - Some(AddColumnLocation::After { - column_name: name.value, - }) - } else { - None - } - } else { - None - }; - Ok(AddColumn { - column_def, - location, - add_if_not_exists, - }) -} - /// Parses a comma separated list of string literals. fn parse_string_option_names(parser: &mut Parser) -> std::result::Result { parser.parse_literal_string() @@ -750,6 +775,75 @@ mod tests { use crate::dialect::GreptimeDbDialect; use crate::parser::ParseOptions; use crate::statements::alter::AlterDatabaseOperation; + use crate::statements::create::{Column, ColumnExtensions}; + + #[test] + fn test_add_column_matches_create_table() { + for (definition, location) in [ + ( + "payload JSON2(max_auto_expanded_paths = 100, service STRING) NOT NULL", + "AFTER ts", + ), + ("payload JSON2(service STRING) NULL", "FIRST"), + ("embedding VECTOR(3) NOT NULL", "AFTER ts"), + ("reading INT DEFAULT 42 COMMENT 'value'", ""), + ("\"select\" STRING NULL", ""), + ] { + let parse = |sql: &str| { + ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .pop() + .unwrap() + }; + let Statement::CreateTable(create) = parse(&format!( + "CREATE TABLE t ({definition}, n INT, ts TIMESTAMP TIME INDEX)" + )) else { + unreachable!() + }; + let Statement::AlterTable(alter) = parse(&format!( + "ALTER TABLE t ADD COLUMN {definition} {location}, ADD COLUMN n INT" + )) else { + unreachable!() + }; + let AlterTableOperation::AddColumns { add_columns } = alter.alter_operation() else { + unreachable!() + }; + assert_eq!(add_columns.len(), 2); + for (expected, actual) in create.columns.iter().zip(add_columns) { + assert_eq!(expected, &actual.column); + } + let Statement::AlterTable(roundtrip) = parse(&alter.to_string()) else { + unreachable!() + }; + assert_eq!(alter, roundtrip); + } + } + + #[test] + fn test_add_column_rejects_unsupported_attributes() { + for definition in [ + "payload STRING FULLTEXT INDEX", + "payload STRING SKIPPING INDEX", + "payload STRING INVERTED INDEX", + "ts TIMESTAMP TIME INDEX", + ] { + let err = ParserContext::create_with_dialect( + &format!("ALTER TABLE t ADD COLUMN {definition}"), + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("ADD COLUMN does not support inline indexes or TIME INDEX"), + "{err}" + ); + } + } #[test] fn test_parse_alter_database() { @@ -830,10 +924,14 @@ mod tests { match alter_operation { AlterTableOperation::AddColumns { add_columns } => { assert_eq!(add_columns.len(), 1); - assert_eq!("tagk_i", add_columns[0].column_def.name.value); - assert_eq!(DataType::String(None), add_columns[0].column_def.data_type); + assert_eq!("tagk_i", add_columns[0].column.column_def.name.value); + assert_eq!( + DataType::String(None), + add_columns[0].column.column_def.data_type + ); assert!( add_columns[0] + .column .column_def .options .iter() @@ -866,10 +964,14 @@ mod tests { assert_matches!(alter_operation, AlterTableOperation::AddColumns { .. }); match alter_operation { AlterTableOperation::AddColumns { add_columns } => { - assert_eq!("tagk_i", add_columns[0].column_def.name.value); - assert_eq!(DataType::String(None), add_columns[0].column_def.data_type); + assert_eq!("tagk_i", add_columns[0].column.column_def.name.value); + assert_eq!( + DataType::String(None), + add_columns[0].column.column_def.data_type + ); assert!( add_columns[0] + .column .column_def .options .iter() @@ -931,7 +1033,7 @@ mod tests { .zip(expecteds) .collect::, ColumnDef))>>() { - assert_eq!(add_column.column_def, expected.1); + assert_eq!(add_column.column.column_def, expected.1); assert_eq!(&expected.0, &add_column.location); } } @@ -962,28 +1064,37 @@ mod tests { AlterTableOperation::AddColumns { add_columns } => { let expected = [ AddColumn { - column_def: ColumnDef { - name: Ident::new("a"), - data_type: DataType::Integer(None), - options: vec![], + column: Column { + column_def: ColumnDef { + name: Ident::new("a"), + data_type: DataType::Integer(None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }, location: None, add_if_not_exists: true, }, AddColumn { - column_def: ColumnDef { - name: Ident::new("b"), - data_type: DataType::String(None), - options: vec![], + column: Column { + column_def: ColumnDef { + name: Ident::new("b"), + data_type: DataType::String(None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }, location: None, add_if_not_exists: false, }, AddColumn { - column_def: ColumnDef { - name: Ident::new("c"), - data_type: DataType::Int(None), - options: vec![], + column: Column { + column_def: ColumnDef { + name: Ident::new("c"), + data_type: DataType::Int(None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }, location: None, add_if_not_exists: true, @@ -1576,9 +1687,8 @@ MODIFY COLUMN attrs JSON2 ( let Statement::AlterTable(alter_table) = statements.remove(0) else { unreachable!() }; - let AlterTableOperation::ModifyColumnType { + let AlterTableOperation::SetJsonSettings { column_name, - target_type, json2_options: Some(options), } = alter_table.alter_operation() else { @@ -1586,7 +1696,6 @@ MODIFY COLUMN attrs JSON2 ( }; 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); @@ -1610,7 +1719,7 @@ MODIFY COLUMN attrs JSON2 ( let Statement::AlterTable(empty) = empty.remove(0) else { unreachable!() }; - let AlterTableOperation::ModifyColumnType { json2_options, .. } = empty.alter_operation() + let AlterTableOperation::SetJsonSettings { json2_options, .. } = empty.alter_operation() else { unreachable!() }; diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index bb73870b45f..c87d0274382 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -41,14 +41,14 @@ use table::requests::{validate_database_option, validate_database_option_value}; use crate::ast::{ColumnDef, Ident, ObjectNamePartExt}; use crate::error::{ self, InvalidColumnOptionSnafu, InvalidDatabaseOptionSnafu, InvalidDatabaseOptionValueSnafu, - InvalidFlowQuerySnafu, InvalidIntervalSnafu, InvalidSqlSnafu, InvalidTimeIndexSnafu, - MissingTimeIndexSnafu, Result, SyntaxSnafu, UnexpectedSnafu, UnsupportedSnafu, + InvalidFlowQuerySnafu, InvalidIntervalSnafu, InvalidTimeIndexSnafu, MissingTimeIndexSnafu, + Result, SyntaxSnafu, UnexpectedSnafu, UnsupportedSnafu, }; use crate::parser::{FLOW, ParserContext}; use crate::parsers::tql_parser; use crate::parsers::utils::{ self, parse_with_options, validate_column_fulltext_create_option, - validate_column_skipping_index_create_option, validate_column_vector_index_create_option, + validate_column_skipping_index_create_option, }; use crate::statements::create::{ Column, ColumnExtensions, CreateDatabase, CreateExternalTable, CreateFlow, CreateTable, @@ -66,7 +66,6 @@ pub const EXPIRE: &str = "EXPIRE"; pub const AFTER: &str = "AFTER"; pub const INVERTED: &str = "INVERTED"; pub const SKIPPING: &str = "SKIPPING"; -pub const VECTOR: &str = "VECTOR"; pub type RawIntervalExpr = String; @@ -686,7 +685,7 @@ impl<'a> ParserContext<'a> { columns: &mut Vec, constraints: &mut Vec, ) -> Result<()> { - let mut column = self.parse_column_def()?; + let mut column = self.parse_column_def(false)?; let mut time_index_opt_idx = None; for (index, opt) in column.options().iter().enumerate() { @@ -769,29 +768,23 @@ impl<'a> ParserContext<'a> { { return Err(ParserError::ParserError(format!( "Cannot use keyword '{}' as column name. Hint: add quotes to the name.", - &name.value + name.value ))); } Ok(name) } - pub fn parse_column_def(&mut self) -> Result { - let name = self.parse_column_name().context(SyntaxSnafu)?; + /// Parses a column definition, optionally allowing unquoted keyword names + /// for compatibility with ALTER TABLE ADD COLUMN. + pub fn parse_column_def(&mut self, allow_keyword_name: bool) -> Result { + let name = if allow_keyword_name { + self.parser.parse_identifier().context(SyntaxSnafu)? + } else { + self.parse_column_name().context(SyntaxSnafu)? + }; let parser = &mut self.parser; - ensure!( - !(name.quote_style.is_none() && - // "ALL_KEYWORDS" are sorted. - ALL_KEYWORDS.binary_search(&name.value.to_uppercase().as_str()).is_ok()), - InvalidSqlSnafu { - msg: format!( - "Cannot use keyword '{}' as column name. Hint: add quotes to the name.", - &name.value - ), - } - ); - let mut extensions = ColumnExtensions::default(); let data_type = @@ -1068,61 +1061,6 @@ impl<'a> ParserContext<'a> { is_index_declared |= true; } - // vector index - if let Token::Word(word) = parser.peek_token().token - && word.value.eq_ignore_ascii_case(VECTOR) - { - parser.next_token(); - // Consume `INDEX` keyword - ensure!( - parser.parse_keyword(Keyword::INDEX), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: "expect INDEX after VECTOR keyword", - } - ); - - ensure!( - column_extensions.vector_index_options.is_none(), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: "duplicated VECTOR INDEX option", - } - ); - - // Check that column is a vector type - let column_type = get_unalias_type(column_type); - let data_type = sql_data_type_to_concrete_data_type(&column_type)?; - ensure!( - matches!(data_type, ConcreteDataType::Vector(_)), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: "VECTOR INDEX only supports Vector type columns", - } - ); - - let options = parser - .parse_options(Keyword::WITH) - .context(error::SyntaxSnafu)? - .into_iter() - .map(parse_option_string) - .collect::>>()?; - - for (key, _) in options.iter() { - ensure!( - validate_column_vector_index_create_option(key), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: format!("invalid VECTOR INDEX option: {key}"), - } - ); - } - - let options = OptionMap::new(options); - column_extensions.vector_index_options = Some(options); - is_index_declared |= true; - } - Ok(is_index_declared) } @@ -2980,11 +2918,16 @@ ENGINE=mito"; #[test] fn test_invalid_column_name() { - let sql = "create table foo(user string, i timestamp time index)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - let err = result.unwrap_err().output_msg(); - assert!(err.contains("Cannot use keyword 'user' as column name")); + for name in ["user", "at", "value", "select"] { + let sql = format!("create table foo({name} string, i timestamp time index)"); + let result = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ); + let err = result.unwrap_err().output_msg(); + assert!(err.contains(&format!("Cannot use keyword '{name}' as column name"))); + } // If column name is quoted, it's valid even same with keyword. let sql = r#" @@ -3526,6 +3469,23 @@ CREATE TABLE log ( } } + #[test] + fn test_reject_vector_index() { + for index in ["VECTOR INDEX", "VECTOR INDEX WITH (metric = 'cosine')"] { + let sql = format!( + "CREATE TABLE vectors (ts TIMESTAMP TIME INDEX, embedding VECTOR(3) {index})" + ); + assert!( + ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err() + ); + } + } + #[test] fn test_parse_interval_cast() { let s = "select '10s'::INTERVAL"; @@ -3534,174 +3494,4 @@ CREATE TABLE log ( .unwrap(); assert_eq!("SELECT '10 seconds'::INTERVAL", &stmts[0].to_string()); } - - #[test] - fn test_parse_create_table_vector_index_options() { - // Test basic vector index - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX, -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap(); - - if let Statement::CreateTable(c) = &result[0] { - c.columns.iter().for_each(|col| { - if col.name().value == "vec" { - assert!( - col.extensions - .vector_index_options - .as_ref() - .unwrap() - .is_empty() - ); - } - }); - } else { - panic!("should be create_table statement"); - } - - // Test vector index with options - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX WITH (metric='cosine', connectivity='32', expansion_add='256', expansion_search='128') -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap(); - - if let Statement::CreateTable(c) = &result[0] { - c.columns.iter().for_each(|col| { - if col.name().value == "vec" { - let options = col.extensions.vector_index_options.as_ref().unwrap(); - assert_eq!(options.len(), 4); - assert_eq!(options.get("metric").unwrap(), "cosine"); - assert_eq!(options.get("connectivity").unwrap(), "32"); - assert_eq!(options.get("expansion_add").unwrap(), "256"); - assert_eq!(options.get("expansion_search").unwrap(), "128"); - } - }); - } else { - panic!("should be create_table statement"); - } - } - - #[test] - fn test_parse_create_table_vector_index_invalid_type() { - // Test vector index on non-vector type (should fail) - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - col INT VECTOR INDEX, -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("VECTOR INDEX only supports Vector type columns") - ); - } - - #[test] - fn test_parse_create_table_vector_index_duplicate() { - // Test duplicate vector index (should fail) - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX VECTOR INDEX, -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("duplicated VECTOR INDEX option") - ); - } - - #[test] - fn test_parse_create_table_vector_index_invalid_option() { - // Test invalid option key (should fail) - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX WITH (metric='l2sq', invalid_option='foo') -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("invalid VECTOR INDEX option") - ); - } - - #[test] - fn test_parse_column_extensions_vector_index() { - // Test vector index on vector type - { - let sql = "VECTOR INDEX WITH (metric = 'l2sq')"; - let dialect = GenericDialect {}; - let mut tokenizer = Tokenizer::new(&dialect, sql); - let tokens = tokenizer.tokenize().unwrap(); - let mut parser = Parser::new(&dialect).with_tokens(tokens); - let name = Ident::new("vec_col"); - let data_type = - DataType::Custom(vec![Ident::new("VECTOR")].into(), vec!["128".to_string()]); - // First, parse the vector type to set vector_options - let mut extensions = ColumnExtensions { - vector_options: Some(OptionMap::from([( - VECTOR_OPT_DIM.to_string(), - "128".to_string(), - )])), - ..Default::default() - }; - - let result = ParserContext::parse_column_extensions( - &mut parser, - &name, - &data_type, - &mut extensions, - ); - assert!(result.is_ok()); - assert!(extensions.vector_index_options.is_some()); - let vi_options = extensions.vector_index_options.unwrap(); - assert_eq!(vi_options.get("metric"), Some("l2sq")); - } - - // Test vector index on non-vector type (should fail) - { - let sql = "VECTOR INDEX"; - let dialect = GenericDialect {}; - let mut tokenizer = Tokenizer::new(&dialect, sql); - let tokens = tokenizer.tokenize().unwrap(); - let mut parser = Parser::new(&dialect).with_tokens(tokens); - let name = Ident::new("num_col"); - let data_type = DataType::Int(None); // Non-vector type - let mut extensions = ColumnExtensions::default(); - let result = ParserContext::parse_column_extensions( - &mut parser, - &name, - &data_type, - &mut extensions, - ); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("VECTOR INDEX only supports Vector type columns") - ); - } - } } diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index 2bf3c9407a6..0adadc1d193 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -25,7 +25,6 @@ use crate::dialect::GreptimeDbDialect; use crate::error::{InvalidSqlSnafu, Result, SyntaxSnafu}; use crate::parsers::create_parser::{INVERTED, SKIPPING}; use crate::statements::create::{Json2Options, JsonTypeHint}; -use crate::statements::transform::type_alias::get_type_by_alias; const JSON2_TYPE_NAME: &str = "JSON2"; const MAX_AUTO_EXPANDED_PATHS: &str = "max_auto_expanded_paths"; @@ -155,7 +154,7 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { } ); let data_type = parser.parse_data_type().context(SyntaxSnafu)?; - let data_type = normalize_json2_type_hint_type(data_type)?; + let data_type = validate_json2_type_hint_type(data_type)?; let mut inverted_index = false; @@ -238,35 +237,23 @@ fn parse_json2_path(parser: &mut Parser<'_>) -> Result> { Ok(path) } -fn normalize_json2_type_hint_type(data_type: DataType) -> Result { - let data_type = get_type_by_alias(&data_type).unwrap_or(data_type); - let normalized = match data_type { - DataType::String(_) | DataType::Text | DataType::Varchar(_) | DataType::Char(_) => { - DataType::String(None) +fn validate_json2_type_hint_type(data_type: DataType) -> Result { + match data_type { + DataType::Int64 + | DataType::UInt64 + | DataType::Float64 + | DataType::String(None) + | DataType::BigInt(None) + | DataType::BigIntUnsigned(None) + | DataType::Double(ExactNumberInfo::None) + | DataType::Boolean => Ok(data_type), + _ => InvalidSqlSnafu { + msg: format!( + "unsupported JSON2 type hint data type: {data_type}; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64" + ), } - DataType::TinyInt(_) - | DataType::SmallInt(_) - | DataType::Int(_) - | DataType::Integer(_) - | DataType::BigInt(_) => DataType::BigInt(None), - DataType::TinyIntUnsigned(_) - | DataType::SmallIntUnsigned(_) - | DataType::IntUnsigned(_) - | DataType::UnsignedInteger - | DataType::BigIntUnsigned(_) => DataType::BigIntUnsigned(None), - DataType::Float(_) | DataType::Real | DataType::Double(_) => { - DataType::Double(ExactNumberInfo::None) - } - DataType::Boolean => DataType::Boolean, - _ => { - return InvalidSqlSnafu { - msg: format!("unsupported JSON2 type hint data type: {data_type}"), - } - .fail(); - } - }; - - Ok(normalized) + .fail(), + } } fn ensure_no_path_conflict(hints: &[JsonTypeHint], path: &[String]) -> Result<()> { @@ -330,7 +317,7 @@ CREATE TABLE traces ( log_json_data JSON2 ( "service.name" STRING INVERTED INDEX, http.method STRING, - status_code INT64, + status_code BIGINT, comment STRING, ), ts TIMESTAMP TIME INDEX, @@ -463,40 +450,81 @@ CREATE TABLE traces ( } #[test] - fn test_parse_json2_type_hint_normalizes_numeric_types() { - let column = parse_json2_column( - r#" -CREATE TABLE traces ( - log_json_data JSON2 ( - tinyint_value TINYINT, - smallint_value SMALLINT, - int_value INT, - integer_value INTEGER, - bigint_value BIGINT, - int64_value INT64, - tinyuint_value TINYINT UNSIGNED, - smalluint_value SMALLINT UNSIGNED, - uint_value INT UNSIGNED, - uint64_value UINT64, - float_value FLOAT, - real_value REAL, - double_value DOUBLE, - float64_value FLOAT64 - ), - ts TIMESTAMP TIME INDEX, -)"#, - ); + fn test_parse_json2_type_hint_supported_types() { + for (sql_type, expected) in [ + ("STRING", DataType::String(None)), + ("BIGINT", DataType::BigInt(None)), + ("BIGINT UNSIGNED", DataType::BigIntUnsigned(None)), + ("DOUBLE", DataType::Double(ExactNumberInfo::None)), + ("BOOLEAN", DataType::Boolean), + ("INT64", DataType::Int64), + ("UINT64", DataType::UInt64), + ("FLOAT64", DataType::Float64), + ] { + for sql_type in [sql_type.to_string(), sql_type.to_lowercase()] { + let column = parse_json2_column(&format!( + "CREATE TABLE traces (j JSON2 (value {sql_type}), ts TIMESTAMP TIME INDEX)" + )); + let options = column.extensions.json2_options.unwrap(); + let settings = options.build_json_settings().unwrap(); + assert_eq!(settings.type_hints().len(), 1); + let hints = options.type_hints; + assert_eq!(hints[0].data_type, expected); + } + } + } - let hints = column.extensions.json2_options.unwrap().type_hints; - assert_eq!(hints.len(), 14); - for hint in hints.iter().take(6) { - assert_eq!(hint.data_type, DataType::BigInt(None)); - } - for hint in hints.iter().skip(6).take(4) { - assert_eq!(hint.data_type, DataType::BigIntUnsigned(None)); - } - for hint in hints.iter().skip(10) { - assert_eq!(hint.data_type, DataType::Double(ExactNumberInfo::None)); + #[test] + fn test_parse_json2_type_hint_rejects_unsupported_types() { + for sql_type in [ + "INT2", + "INT4", + "INT8", + "INT16", + "INT32", + "TINYINT", + "SMALLINT", + "INT", + "INTEGER", + "UINT8", + "UINT16", + "UINT32", + "TINYINT UNSIGNED", + "SMALLINT UNSIGNED", + "INT UNSIGNED", + "FLOAT", + "REAL", + "FLOAT4", + "FLOAT8", + "FLOAT32", + "BOOL", + "TEXT", + "VARCHAR(10)", + "CHAR(10)", + "TIMESTAMP", + "DECIMAL(10, 2)", + "STRING(10)", + "BIGINT(10)", + ] { + for sql in [ + format!( + "CREATE TABLE traces (j JSON2 (value {sql_type}), ts TIMESTAMP TIME INDEX)" + ), + format!("ALTER TABLE traces MODIFY COLUMN j JSON2 (value {sql_type})"), + ] { + let err = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap_err(); + assert!( + err.to_string().contains( + "supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64" + ), + "{sql}: {err}" + ); + } } } @@ -504,7 +532,7 @@ CREATE TABLE traces ( fn test_parse_json2_type_hint_rejects_default() { for default in ["-5", "abs(-1)"] { let sql = format!( - "CREATE TABLE traces (log_json_data JSON2 (status_code INT64 DEFAULT {default}), ts TIMESTAMP TIME INDEX)" + "CREATE TABLE traces (log_json_data JSON2 (status_code BIGINT DEFAULT {default}), ts TIMESTAMP TIME INDEX)" ); let err = ParserContext::create_with_dialect( &sql, @@ -521,7 +549,7 @@ CREATE TABLE traces ( let result = ParserContext::create_with_dialect( r#" CREATE TABLE traces ( - log_json_data JSON2 (a.b STRING, a.b INT64), + log_json_data JSON2 (a.b STRING, a.b BIGINT), ts TIMESTAMP TIME INDEX, )"#, &GreptimeDbDialect {}, @@ -537,7 +565,7 @@ CREATE TABLE traces ( let result = ParserContext::create_with_dialect( r#" CREATE TABLE traces ( - log_json_data JSON2 (a STRING, a.b INT64), + log_json_data JSON2 (a STRING, a.b BIGINT), ts TIMESTAMP TIME INDEX, )"#, &GreptimeDbDialect {}, diff --git a/src/sql/src/parsers/utils.rs b/src/sql/src/parsers/utils.rs index e91a2344643..afba437045c 100644 --- a/src/sql/src/parsers/utils.rs +++ b/src/sql/src/parsers/utils.rs @@ -32,9 +32,6 @@ use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, - COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, - COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, }; use snafu::{ResultExt, ensure}; use sqlparser::dialect::Dialect; @@ -378,17 +375,6 @@ pub fn validate_column_skipping_index_create_option(key: &str) -> bool { .contains(&key) } -pub fn validate_column_vector_index_create_option(key: &str) -> bool { - [ - COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, - COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, - COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, - ] - .contains(&key) -} - /// Convert an [`IntervalMonthDayNano`] to a [`Duration`]. #[cfg(feature = "enterprise")] pub fn convert_month_day_nano_to_duration( diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 808e30f842c..c2c97c34567 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -36,7 +36,10 @@ pub mod truncate; use std::sync::Arc; use api::helper::ColumnDataTypeWrapper; -use api::v1::SemanticType; +use api::v1::{ColumnOptions, SemanticType}; +use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, +}; use common_sql::default_constraint::parse_column_default_constraint; use common_time::timezone::Timezone; use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; @@ -56,7 +59,7 @@ use crate::ast::{ use crate::error::{ self, ConvertToGrpcDataTypeSnafu, ConvertValueSnafu, Result, SerializeColumnDefaultConstraintSnafu, SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, - SetVectorIndexOptionSnafu, SqlCommonSnafu, + SqlCommonSnafu, }; use crate::statements::create::Column; pub use crate::statements::option_map::OptionMap; @@ -146,40 +149,36 @@ pub fn column_to_schema( .context(SetSkippingIndexOptionSnafu)?; } - if let Some(options) = column.extensions.build_vector_index_options()? { - column_schema = column_schema - .with_vector_index_options(&options) - .context(SetVectorIndexOptionSnafu)?; - } - column_schema.set_inverted_index(column.extensions.inverted_index_options.is_some()); - let is_json2_column = if let SqlDataType::Custom(object_name, _) = column.data_type() { - object_name - .0 - .first() - .map(|x| x.to_string_unquoted().eq_ignore_ascii_case(JSON2_TYPE_NAME)) - .unwrap_or_default() - } else { - false - }; - if is_json2_column { - let settings = column - .extensions - .build_json_settings()? - .unwrap_or_else(JsonSettings::new_v2); - let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))); - column_schema.with_extension_type(&extension); - } + set_json2_extension(&mut column_schema, column)?; Ok(column_schema) } -/// Convert `ColumnDef` in sqlparser to `ColumnDef` in gRPC proto. -pub fn sql_column_def_to_grpc_column_def( - col: &ColumnDef, +/// Attaches extension metadata and settings to JSON2 columns. +fn set_json2_extension(column_schema: &mut ColumnSchema, column: &Column) -> Result<()> { + if !column_schema.data_type.is_json2() { + return Ok(()); + } + + let settings = column + .extensions + .build_json_settings()? + .unwrap_or_else(JsonSettings::new_v2); + if let Some(extension) = json2_extension(column.data_type(), settings) { + column_schema.with_extension_type(&extension); + } + + Ok(()) +} + +/// Converts a SQL [`Column`] to a `ColumnDef` in gRPC proto. +pub fn sql_col_to_grpc_col_def( + column: &Column, timezone: Option<&Timezone>, ) -> Result { + let col = &column.column_def; let name = col.name.value.clone(); let data_type = sql_data_type_to_concrete_data_type(&col.data_type)?; @@ -210,6 +209,24 @@ pub fn sql_column_def_to_grpc_column_def( SemanticType::Field }; + let settings = column + .extensions + .build_json_settings()? + .unwrap_or_else(JsonSettings::new_v2); + let options = json2_extension(&col.data_type, settings).map(|extension| { + let mut options = ColumnOptions::default(); + options.options.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + Json2ExtensionType::NAME.to_string(), + ); + if let Some(metadata) = extension.serialize_metadata() { + options + .options + .insert(EXTENSION_TYPE_METADATA_KEY.to_string(), metadata); + } + options + }); + Ok(api::v1::ColumnDef { name, data_type: datatype as i32, @@ -218,14 +235,32 @@ pub fn sql_column_def_to_grpc_column_def( semantic_type: semantic_type as _, comment: String::new(), datatype_extension: datatype_ext, - options: None, + options, }) } +fn json2_extension(data_type: &SqlDataType, settings: JsonSettings) -> Option { + let SqlDataType::Custom(name, _) = data_type else { + return None; + }; + if !name.0.first().is_some_and(|name| { + name.to_string_unquoted() + .eq_ignore_ascii_case(JSON2_TYPE_NAME) + }) { + return None; + } + + Some(Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings, + )))) +} + pub fn sql_data_type_to_concrete_data_type(data_type: &SqlDataType) -> Result { match data_type { SqlDataType::BigInt(_) | SqlDataType::Int64 => Ok(ConcreteDataType::int64_datatype()), - SqlDataType::BigIntUnsigned(_) => Ok(ConcreteDataType::uint64_datatype()), + SqlDataType::BigIntUnsigned(_) | SqlDataType::UInt64 => { + Ok(ConcreteDataType::uint64_datatype()) + } SqlDataType::Int(_) | SqlDataType::Integer(_) => Ok(ConcreteDataType::int32_datatype()), SqlDataType::IntUnsigned(_) | SqlDataType::UnsignedInteger => { Ok(ConcreteDataType::uint32_datatype()) @@ -422,6 +457,7 @@ mod tests { SqlDataType::BigIntUnsigned(None), ConcreteDataType::uint64_datatype(), ); + check_type(SqlDataType::UInt64, ConcreteDataType::uint64_datatype()); check_type( SqlDataType::IntUnsigned(None), ConcreteDataType::uint32_datatype(), @@ -456,15 +492,18 @@ mod tests { } #[test] - pub fn test_sql_column_def_to_grpc_column_def() { + pub fn test_sql_col_to_grpc_col_def() { // test basic - let column_def = ColumnDef { - name: "col".into(), - data_type: SqlDataType::Double(ExactNumberInfo::None), - options: vec![], + let column = Column { + column_def: ColumnDef { + name: "col".into(), + data_type: SqlDataType::Double(ExactNumberInfo::None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }; - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert_eq!("col", grpc_column_def.name); assert!(grpc_column_def.is_nullable); // nullable when options are empty @@ -473,56 +512,65 @@ mod tests { assert_eq!(grpc_column_def.semantic_type, SemanticType::Field as i32); // test not null - let column_def = ColumnDef { - name: "col".into(), - data_type: SqlDataType::Double(ExactNumberInfo::None), - options: vec![ColumnOptionDef { - name: None, - option: ColumnOption::NotNull, - }], + let column = Column { + column_def: ColumnDef { + name: "col".into(), + data_type: SqlDataType::Double(ExactNumberInfo::None), + options: vec![ColumnOptionDef { + name: None, + option: ColumnOption::NotNull, + }], + }, + extensions: ColumnExtensions::default(), }; - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert!(!grpc_column_def.is_nullable); // test primary key - let column_def = ColumnDef { - name: "col".into(), - data_type: SqlDataType::Double(ExactNumberInfo::None), - options: vec![ColumnOptionDef { - name: None, - option: ColumnOption::PrimaryKey(PrimaryKeyConstraint { + let column = Column { + column_def: ColumnDef { + name: "col".into(), + data_type: SqlDataType::Double(ExactNumberInfo::None), + options: vec![ColumnOptionDef { name: None, - index_name: None, - index_type: None, - columns: vec![], - index_options: vec![], - characteristics: None, - }), - }], + option: ColumnOption::PrimaryKey(PrimaryKeyConstraint { + name: None, + index_name: None, + index_type: None, + columns: vec![], + index_options: vec![], + characteristics: None, + }), + }], + }, + extensions: ColumnExtensions::default(), }; - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert_eq!(grpc_column_def.semantic_type, SemanticType::Tag as i32); } #[test] - pub fn test_sql_column_def_to_grpc_column_def_with_timezone() { - let column_def = ColumnDef { - name: "col".into(), - // MILLISECOND - data_type: SqlDataType::Timestamp(Some(3), TimezoneInfo::None), - options: vec![ColumnOptionDef { - name: None, - option: ColumnOption::Default(Expr::Value( - SqlValue::SingleQuotedString("2024-01-30T00:01:01".to_string()).into(), - )), - }], + pub fn test_sql_col_to_grpc_col_def_with_timezone() { + let column = Column { + column_def: ColumnDef { + name: "col".into(), + // MILLISECOND + data_type: SqlDataType::Timestamp(Some(3), TimezoneInfo::None), + options: vec![ColumnOptionDef { + name: None, + option: ColumnOption::Default(Expr::Value( + SqlValue::SingleQuotedString("2024-01-30T00:01:01".to_string()).into(), + )), + }], + }, + extensions: ColumnExtensions::default(), }; // with timezone "Asia/Shanghai" - let grpc_column_def = sql_column_def_to_grpc_column_def( - &column_def, + let grpc_column_def = sql_col_to_grpc_col_def( + &column, Some(&Timezone::from_tz_string("Asia/Shanghai").unwrap()), ) .unwrap(); @@ -542,7 +590,7 @@ mod tests { ); // without timezone - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert_eq!("col", grpc_column_def.name); assert!(grpc_column_def.is_nullable); // nullable when options are empty assert_eq!( @@ -778,74 +826,4 @@ mod tests { assert_eq!(fulltext_options.analyzer, FulltextAnalyzer::English); assert!(fulltext_options.case_sensitive); } - - #[test] - fn test_column_to_schema_with_vector_index() { - use datatypes::schema::{VectorDistanceMetric, VectorIndexEngineType}; - - // Test with custom metric and parameters - let column = Column { - column_def: ColumnDef { - name: "embedding".into(), - data_type: SqlDataType::Custom( - vec![Ident::new(VECTOR_TYPE_NAME)].into(), - vec!["128".to_string()], - ), - options: vec![], - }, - extensions: ColumnExtensions { - vector_index_options: Some(OptionMap::from([ - ("metric".to_string(), "cosine".to_string()), - ("connectivity".to_string(), "32".to_string()), - ("expansion_add".to_string(), "200".to_string()), - ("expansion_search".to_string(), "100".to_string()), - ])), - ..Default::default() - }, - }; - - let column_schema = column_to_schema(&column, "ts", None).unwrap(); - assert_eq!("embedding", column_schema.name); - assert!(column_schema.is_vector_indexed()); - - let vector_options = column_schema.vector_index_options().unwrap().unwrap(); - assert_eq!(vector_options.engine, VectorIndexEngineType::Usearch); - assert_eq!(vector_options.metric, VectorDistanceMetric::Cosine); - assert_eq!(vector_options.connectivity, 32); - assert_eq!(vector_options.expansion_add, 200); - assert_eq!(vector_options.expansion_search, 100); - } - - #[test] - fn test_column_to_schema_with_vector_index_defaults() { - use datatypes::schema::{VectorDistanceMetric, VectorIndexEngineType}; - - // Test with default values (empty options map) - let column = Column { - column_def: ColumnDef { - name: "vec".into(), - data_type: SqlDataType::Custom( - vec![Ident::new(VECTOR_TYPE_NAME)].into(), - vec!["64".to_string()], - ), - options: vec![], - }, - extensions: ColumnExtensions { - vector_index_options: Some(OptionMap::default()), - ..Default::default() - }, - }; - - let column_schema = column_to_schema(&column, "ts", None).unwrap(); - assert_eq!("vec", column_schema.name); - assert!(column_schema.is_vector_indexed()); - - let vector_options = column_schema.vector_index_options().unwrap().unwrap(); - // Verify defaults - assert_eq!(vector_options.engine, VectorIndexEngineType::Usearch); - assert_eq!(vector_options.metric, VectorDistanceMetric::L2sq); - assert_eq!(vector_options.connectivity, 16); - assert_eq!(vector_options.expansion_add, 128); - assert_eq!(vector_options.expansion_search, 64); - } } diff --git a/src/sql/src/statements/alter.rs b/src/sql/src/statements/alter.rs index 2f98d8737ae..781c65ab528 100644 --- a/src/sql/src/statements/alter.rs +++ b/src/sql/src/statements/alter.rs @@ -22,11 +22,11 @@ use common_query::AddColumnLocation; use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use itertools::Itertools; use serde::Serialize; -use sqlparser::ast::{ColumnDef, DataType, Expr, Ident, ObjectName, TableConstraint}; +use sqlparser::ast::{DataType, Expr, Ident, ObjectName, TableConstraint}; use sqlparser_derive::{Visit, VisitMut}; use crate::statements::OptionMap; -use crate::statements::create::{Json2Options, Partitions}; +use crate::statements::create::{Column, Json2Options, Partitions}; #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct AlterTable { @@ -88,6 +88,11 @@ pub enum AlterTableOperation { target_type: DataType, json2_options: Option, }, + /// `MODIFY JSON2 [json2_options]` + SetJsonSettings { + column_name: Ident, + json2_options: Option, + }, /// `SET
=
` SetTableOptions { options: Vec, @@ -224,18 +229,18 @@ pub enum UnsetIndexOperation { #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct AddColumn { - pub column_def: ColumnDef, + pub column: Column, pub location: Option, pub add_if_not_exists: bool, } impl Display for AddColumn { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.column)?; if let Some(location) = &self.location { - write!(f, "{} {location}", self.column_def) - } else { - write!(f, "{}", self.column_def) + write!(f, " {location}")?; } + Ok(()) } } @@ -265,6 +270,16 @@ impl Display for AlterTableOperation { } Ok(()) } + AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + } => { + write!(f, r#"MODIFY COLUMN {column_name} JSON2"#)?; + if let Some(options) = json2_options { + write!(f, "{options}")?; + } + Ok(()) + } AlterTableOperation::SetTableOptions { options } => { let kvs = options .iter() @@ -433,8 +448,10 @@ impl Display for AlterDatabaseOperation { mod tests { use std::assert_matches; + use super::AlterTableOperation; use crate::dialect::GreptimeDbDialect; use crate::parser::{ParseOptions, ParserContext}; + use crate::statements::create::Json2Options; use crate::statements::statement::Statement; #[test] @@ -503,7 +520,7 @@ ALTER TABLE monitor ADD COLUMN app STRING DEFAULT 'shop' PRIMARY KEY, ADD COLUMN } let sql = r"alter table monitor modify column load_15 string;"; - let stmts = + let mut stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) .unwrap(); assert_eq!(1, stmts.len()); @@ -523,6 +540,25 @@ ALTER TABLE monitor MODIFY COLUMN load_15 STRING"#, } } + let Statement::AlterTable(alter_table) = &mut stmts[0] else { + unreachable!(); + }; + let AlterTableOperation::ModifyColumnType { json2_options, .. } = + alter_table.alter_operation_mut() + else { + unreachable!(); + }; + *json2_options = Some(Json2Options { + max_auto_expanded_paths: Some(1), + type_hints: vec![], + }); + assert_eq!( + r#"ALTER TABLE monitor MODIFY COLUMN load_15 STRING( + max_auto_expanded_paths = 1 + )"#, + alter_table.to_string() + ); + let sql = r"alter table monitor drop column load_15;"; let stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) diff --git a/src/sql/src/statements/copy.rs b/src/sql/src/statements/copy.rs index 7a59175538c..f47f26ec61a 100644 --- a/src/sql/src/statements/copy.rs +++ b/src/sql/src/statements/copy.rs @@ -46,7 +46,7 @@ pub struct CopyQueryTo { impl Display for CopyQueryTo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "COPY ({}) TO {}", &self.query, &self.arg.location)?; + write!(f, "COPY ({}) TO {}", self.query, self.arg.location)?; if !self.arg.with.is_empty() { let options = self.arg.with.kv_pairs(); write!(f, " WITH ({})", options.join(", "))?; @@ -70,11 +70,11 @@ impl Display for CopyTable { write!(f, "COPY ")?; let (with, connection) = match self { CopyTable::To(args) => { - write!(f, "{} TO {}", &args.table_name, &args.location)?; + write!(f, "{} TO {}", args.table_name, args.location)?; (&args.with, &args.connection) } CopyTable::From(args) => { - write!(f, "{} FROM {}", &args.table_name, &args.location)?; + write!(f, "{} FROM {}", args.table_name, args.location)?; (&args.with, &args.connection) } }; @@ -101,11 +101,11 @@ impl Display for CopyDatabase { write!(f, "COPY DATABASE ")?; let (with, connection) = match self { CopyDatabase::To(args) => { - write!(f, "{} TO {}", &args.database_name, &args.location)?; + write!(f, "{} TO {}", args.database_name, args.location)?; (&args.with, &args.connection) } CopyDatabase::From(args) => { - write!(f, "{} FROM {}", &args.database_name, &args.location)?; + write!(f, "{} FROM {}", args.database_name, args.location)?; (&args.with, &args.connection) } }; diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 46b5b07c9f7..919c1cd31f9 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -13,22 +13,19 @@ // limitations under the License. use std::collections::HashMap; -use std::fmt::{Display, Formatter}; +use std::fmt::{Display, Formatter, Write}; use common_catalog::consts::FILE_ENGINE; use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings}; use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{ - FulltextOptions, SkippingIndexOptions, VectorDistanceMetric, VectorIndexEngineType, - VectorIndexOptions, -}; +use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use itertools::Itertools; use serde::Serialize; use snafu::ResultExt; use sqlparser::ast::{ColumnOptionDef, DataType, Expr}; use sqlparser_derive::{Visit, VisitMut}; -use crate::ast::{ColumnDef, Ident, ObjectName, Value as SqlValue}; +use crate::ast::{ColumnDef, Ident, ObjectName}; use crate::error::{ InvalidFlowQuerySnafu, InvalidSqlSnafu, Result, SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, @@ -71,6 +68,36 @@ fn format_table_constraint(constraints: &[TableConstraint]) -> String { constraints.iter().map(|c| format_indent!(c)).join(LINE_SEP) } +fn write_list( + f: &mut dyn Write, + values: &[T], + separator: &str, + prefix: &str, +) -> std::fmt::Result { + for (index, value) in values.iter().enumerate() { + if index > 0 { + f.write_str(separator)?; + } + write!(f, "{prefix}{value}")?; + } + Ok(()) +} + +// sqlparser's Custom Display joins all modifiers into an intermediate String. +fn write_data_type(f: &mut dyn Write, data_type: &DataType) -> std::fmt::Result { + if let DataType::Custom(name, modifiers) = data_type { + write!(f, "{name}")?; + if !modifiers.is_empty() { + f.write_char('(')?; + write_list(f, modifiers, COMMA_SEP, "")?; + f.write_char(')')?; + } + Ok(()) + } else { + write!(f, "{data_type}") + } +} + /// Table constraint for create table statement. #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] pub enum TableConstraint { @@ -84,7 +111,9 @@ impl Display for TableConstraint { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { TableConstraint::PrimaryKey { columns } => { - write!(f, "PRIMARY KEY ({})", format_list_comma!(columns)) + f.write_str("PRIMARY KEY (")?; + write_list(f, columns, COMMA_SEP, "")?; + f.write_str(")") } TableConstraint::TimeIndex { column } => { write!(f, "TIME INDEX ({})", column) @@ -108,7 +137,7 @@ pub struct CreateTable { pub partitions: Option, } -/// Column definition in `CREATE TABLE` statement. +/// Column definition in `CREATE TABLE` and `ALTER TABLE ADD COLUMN`. #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] pub struct Column { /// `ColumnDef` from `sqlparser::ast` @@ -131,8 +160,6 @@ pub struct ColumnExtensions { /// /// Inverted index doesn't have options at present. There won't be any options in that map. pub inverted_index_options: Option, - /// Vector index options for HNSW-based vector similarity search. - pub vector_index_options: Option, /// JSON2-specific column options. pub json2_options: Option, } @@ -148,12 +175,54 @@ pub struct Json2Options { impl Display for Json2Options { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - let mut options = Vec::with_capacity(self.type_hints.len() + 1); + f.write_str("(\n ")?; if let Some(max) = self.max_auto_expanded_paths { - options.push(format!("max_auto_expanded_paths = {max}")); + write!(f, "max_auto_expanded_paths = {max}")?; } - options.extend(self.type_hints.iter().map(format_json_type_hint)); - write!(f, "(\n {}\n )", options.iter().join(",\n ")) + for (index, hint) in self.type_hints.iter().enumerate() { + if index > 0 || self.max_auto_expanded_paths.is_some() { + f.write_str(",\n ")?; + } + for (index, segment) in hint.path.iter().enumerate() { + if index > 0 { + f.write_char('.')?; + } + f.write_char('"')?; + for part in segment.split_inclusive('"') { + f.write_str(part)?; + if part.ends_with('"') { + f.write_char('"')?; + } + } + f.write_char('"')?; + } + f.write_char(' ')?; + write_data_type(f, &hint.data_type)?; + if hint.inverted_index { + f.write_str(" INVERTED INDEX")?; + } + } + f.write_str("\n )") + } +} + +impl Json2Options { + pub fn build_json_settings(&self) -> Result { + let type_hints = self + .type_hints + .iter() + .map(|hint| { + Ok(datatypes::json::JsonTypeHint { + path: hint.path.clone(), + data_type: sql_data_type_to_concrete_data_type(&hint.data_type)?, + inverted_index: hint.inverted_index, + }) + }) + .collect::>>()?; + let max_auto_expanded_paths = self + .max_auto_expanded_paths + .or(Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS)); + JsonSettings::try_new(type_hints, max_auto_expanded_paths).map_err(Into::into) } } @@ -188,14 +257,21 @@ impl Column { impl Display for Column { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - if let Some(vector_options) = &self.extensions.vector_options + self.write_sql(f, false) + } +} + +impl Column { + fn write_sql(&self, f: &mut dyn Write, complete: bool) -> std::fmt::Result { + if !complete + && let Some(vector_options) = &self.extensions.vector_options && let Some(dim) = vector_options.get(VECTOR_OPT_DIM) { write!(f, "{} VECTOR({})", self.column_def.name, dim)?; - return Ok(()); + } else { + write!(f, "{} ", self.column_def.name)?; + write_data_type(f, &self.column_def.data_type)?; } - - write!(f, "{} {}", self.column_def.name, self.column_def.data_type)?; if let Some(options) = &self.extensions.json2_options { write!(f, "{options}")?; } @@ -203,41 +279,21 @@ impl Display for Column { write!(f, " {option}")?; } - if let Some(fulltext_options) = &self.extensions.fulltext_index_options { - if !fulltext_options.is_empty() { - let options = fulltext_options.kv_pairs(); - write!(f, " FULLTEXT INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " FULLTEXT INDEX")?; + for (kind, options) in [ + ("FULLTEXT", &self.extensions.fulltext_index_options), + ("SKIPPING", &self.extensions.skipping_index_options), + ("INVERTED", &self.extensions.inverted_index_options), + ] { + if let Some(options) = options { + write!(f, " {kind} INDEX")?; + if !options.is_empty() { + f.write_str(" WITH(")?; + options.write_sql(f, COMMA_SEP, "", complete)?; + f.write_str(")")?; + } } } - if let Some(skipping_index_options) = &self.extensions.skipping_index_options { - if !skipping_index_options.is_empty() { - let options = skipping_index_options.kv_pairs(); - write!(f, " SKIPPING INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " SKIPPING INDEX")?; - } - } - - if let Some(inverted_index_options) = &self.extensions.inverted_index_options { - if !inverted_index_options.is_empty() { - let options = inverted_index_options.kv_pairs(); - write!(f, " INVERTED INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " INVERTED INDEX")?; - } - } - - if let Some(vector_index_options) = &self.extensions.vector_index_options { - if !vector_index_options.is_empty() { - let options = vector_index_options.kv_pairs(); - write!(f, " VECTOR INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " VECTOR INDEX")?; - } - } Ok(()) } } @@ -263,112 +319,12 @@ impl ColumnExtensions { )) } - pub fn build_vector_index_options(&self) -> Result> { - let Some(options) = self.vector_index_options.as_ref() else { - return Ok(None); - }; - - let options_map: HashMap = options.clone().into_map(); - let mut result = VectorIndexOptions::default(); - - if let Some(s) = options_map.get("engine") { - result.engine = s.parse::().map_err(|e| { - InvalidSqlSnafu { - msg: format!("invalid VECTOR INDEX engine: {e}"), - } - .build() - })?; - } - - if let Some(s) = options_map.get("metric") { - result.metric = s.parse::().map_err(|e| { - InvalidSqlSnafu { - msg: format!("invalid VECTOR INDEX metric: {e}"), - } - .build() - })?; - } - - if let Some(s) = options_map.get("connectivity") { - let value = s.parse::().map_err(|_| { - InvalidSqlSnafu { - msg: format!( - "invalid VECTOR INDEX connectivity: {s}, expected positive integer" - ), - } - .build() - })?; - if !(2..=2048).contains(&value) { - return InvalidSqlSnafu { - msg: "VECTOR INDEX connectivity must be in the range [2, 2048].".to_string(), - } - .fail(); - } - result.connectivity = value; - } - - if let Some(s) = options_map.get("expansion_add") { - let value = s.parse::().map_err(|_| { - InvalidSqlSnafu { - msg: format!( - "invalid VECTOR INDEX expansion_add: {s}, expected positive integer" - ), - } - .build() - })?; - if value == 0 { - return InvalidSqlSnafu { - msg: "VECTOR INDEX expansion_add must be greater than 0".to_string(), - } - .fail(); - } - result.expansion_add = value; - } - - if let Some(s) = options_map.get("expansion_search") { - let value = s.parse::().map_err(|_| { - InvalidSqlSnafu { - msg: format!( - "invalid VECTOR INDEX expansion_search: {s}, expected positive integer" - ), - } - .build() - })?; - if value == 0 { - return InvalidSqlSnafu { - msg: "VECTOR INDEX expansion_search must be greater than 0".to_string(), - } - .fail(); - } - result.expansion_search = value; - } - - Ok(Some(result)) - } - pub fn build_json_settings(&self) -> Result> { let Some(options) = &self.json2_options else { return Ok(None); }; - let type_hints = options - .type_hints - .iter() - .map(|hint| { - Ok(datatypes::json::JsonTypeHint { - path: hint.path.clone(), - data_type: json_type_hint_concrete_data_type(&hint.data_type)?, - inverted_index: hint.inverted_index, - }) - }) - .collect::>>()?; - let settings = JsonSettings::try_new( - type_hints, - options - .max_auto_expanded_paths - .or(Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS)), - )?; - Ok(Some(settings)) + options.build_json_settings().map(Some) } pub fn set_json_settings(&mut self, settings: JsonSettings) -> Result<()> { @@ -393,28 +349,13 @@ impl ColumnExtensions { } } -fn json_type_hint_concrete_data_type(data_type: &DataType) -> Result { - let data_type = sql_data_type_to_concrete_data_type(data_type)?; - normalize_json_type_hint_concrete_data_type(&data_type) -} - -fn normalize_json_type_hint_concrete_data_type( - data_type: &ConcreteDataType, -) -> Result { - let normalized = match data_type { - ConcreteDataType::String(_) => ConcreteDataType::string_datatype(), - ConcreteDataType::Int8(_) - | ConcreteDataType::Int16(_) - | ConcreteDataType::Int32(_) - | ConcreteDataType::Int64(_) => ConcreteDataType::int64_datatype(), - ConcreteDataType::UInt8(_) - | ConcreteDataType::UInt16(_) - | ConcreteDataType::UInt32(_) - | ConcreteDataType::UInt64(_) => ConcreteDataType::uint64_datatype(), - ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_) => { - ConcreteDataType::float64_datatype() - } - ConcreteDataType::Boolean(_) => ConcreteDataType::boolean_datatype(), +fn json_type_hint_sql_data_type(data_type: &ConcreteDataType) -> Result { + let sql_type = match data_type { + ConcreteDataType::String(_) => DataType::String(None), + ConcreteDataType::Int64(_) => DataType::BigInt(None), + ConcreteDataType::UInt64(_) => DataType::BigIntUnsigned(None), + ConcreteDataType::Float64(_) => DataType::Double(sqlparser::ast::ExactNumberInfo::None), + ConcreteDataType::Boolean(_) => DataType::Boolean, _ => { return InvalidSqlSnafu { msg: format!("unsupported JSON2 type hint data type: {data_type}"), @@ -422,40 +363,9 @@ fn normalize_json_type_hint_concrete_data_type( .fail(); } }; - Ok(normalized) -} - -fn json_type_hint_sql_data_type(data_type: &ConcreteDataType) -> Result { - let data_type = normalize_json_type_hint_concrete_data_type(data_type)?; - let sql_type = match data_type { - ConcreteDataType::String(_) => DataType::String(None), - ConcreteDataType::Int64(_) => DataType::BigInt(None), - ConcreteDataType::UInt64(_) => DataType::BigIntUnsigned(None), - ConcreteDataType::Float64(_) => DataType::Double(sqlparser::ast::ExactNumberInfo::None), - ConcreteDataType::Boolean(_) => DataType::Boolean, - _ => unreachable!("JSON2 type hint data type should have been normalized"), - }; Ok(sql_type) } -fn format_json_type_hint(hint: &JsonTypeHint) -> String { - let path = hint - .path - .iter() - .map(|segment| format_json_path_segment(segment)) - .join("."); - let inverted_index = if hint.inverted_index { - " INVERTED INDEX" - } else { - "" - }; - format!("{} {}{}", path, hint.data_type, inverted_index) -} - -fn format_json_path_segment(segment: &str) -> String { - format!("\"{}\"", segment.replace('"', "\"\"")) -} - /// Partition on columns or values. /// /// - `column_list` is the list of columns in `PARTITION ON COLUMNS` clause. @@ -477,32 +387,14 @@ impl Partitions { } } -#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut)] -pub struct PartitionEntry { - pub name: Ident, - pub value_list: Vec, -} - -impl Display for PartitionEntry { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!( - f, - "PARTITION {} VALUES LESS THAN ({})", - self.name, - format_list_comma!(self.value_list), - ) - } -} - impl Display for Partitions { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { if !self.column_list.is_empty() { - write!( - f, - "PARTITION ON COLUMNS ({}) (\n{}\n)", - format_list_comma!(self.column_list), - format_list_indent!(self.exprs), - )?; + f.write_str("PARTITION ON COLUMNS (")?; + write_list(f, &self.column_list, COMMA_SEP, "")?; + f.write_str(") (\n")?; + write_list(f, &self.exprs, LINE_SEP, " ")?; + f.write_str("\n)")?; } Ok(()) } @@ -510,6 +402,53 @@ impl Display for Partitions { impl Display for CreateTable { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.write_sql(f, false) + } +} + +impl CreateTable { + /// Debits the SQL size of a rewritten AST, including unredacted and non-scalar + /// options. Returns false on exhaustion or a shape hidden by normal Display. + /// No SQL or secret value is returned. Delegated sqlparser formatters retain + /// their own allocation behavior; this is not a general AST memory limit. + pub fn consume_sql_size_budget(&self, remaining: &mut usize) -> bool { + struct Budget<'a>(&'a mut usize); + impl Write for Budget<'_> { + fn write_str(&mut self, value: &str) -> std::fmt::Result { + *self.0 = self.0.checked_sub(value.len()).ok_or(std::fmt::Error)?; + Ok(()) + } + } + if self + .partitions + .as_ref() + .is_some_and(|p| p.column_list.is_empty() && !p.exprs.is_empty()) + { + return false; + } + for column in &self.columns { + if let Some(options) = &column.extensions.vector_options { + let DataType::Custom(name, tokens) = &column.column_def.data_type else { + return false; + }; + if options.len() != 1 + || name.0.len() != 1 + || tokens.len() != 1 + || !name.0[0] + .as_ident() + .is_some_and(|name| name.value.eq_ignore_ascii_case("VECTOR")) + || !tokens[0].parse::().ok().is_some_and(|dim| { + options.get(VECTOR_OPT_DIM) == Some(dim.to_string().as_str()) + }) + { + return false; + } + } + } + self.write_sql(&mut Budget(remaining), true).is_ok() + } + + fn write_sql(&self, f: &mut dyn Write, complete: bool) -> std::fmt::Result { write!(f, "CREATE ")?; if self.engine == FILE_ENGINE { write!(f, "EXTERNAL ")?; @@ -518,17 +457,26 @@ impl Display for CreateTable { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - writeln!(f, "{} (", &self.name)?; - writeln!(f, "{},", format_list_indent!(self.columns))?; - writeln!(f, "{}", format_table_constraint(&self.constraints))?; + writeln!(f, "{} (", self.name)?; + for (index, column) in self.columns.iter().enumerate() { + if index > 0 { + f.write_str(LINE_SEP)?; + } + f.write_str(" ")?; + column.write_sql(f, complete)?; + } + f.write_str(",\n")?; + write_list(f, &self.constraints, LINE_SEP, " ")?; + f.write_char('\n')?; writeln!(f, ")")?; if let Some(partitions) = &self.partitions { writeln!(f, "{partitions}")?; } - writeln!(f, "ENGINE={}", &self.engine)?; + writeln!(f, "ENGINE={}", self.engine)?; if !self.options.is_empty() { - let options = self.options.kv_pairs(); - write!(f, "WITH(\n{}\n)", format_list_indent!(options))?; + f.write_str("WITH(\n")?; + self.options.write_sql(f, LINE_SEP, " ", complete)?; + f.write_str("\n)")?; } Ok(()) } @@ -559,7 +507,7 @@ impl Display for CreateDatabase { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - write!(f, "{}", &self.name)?; + write!(f, "{}", self.name)?; if !self.options.is_empty() { let options = self.options.kv_pairs(); write!(f, "\nWITH(\n{}\n)", format_list_indent!(options))?; @@ -586,11 +534,11 @@ impl Display for CreateExternalTable { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - writeln!(f, "{} (", &self.name)?; + writeln!(f, "{} (", self.name)?; writeln!(f, "{},", format_list_indent!(self.columns))?; writeln!(f, "{}", format_table_constraint(&self.constraints))?; writeln!(f, ")")?; - writeln!(f, "ENGINE={}", &self.engine)?; + writeln!(f, "ENGINE={}", self.engine)?; if !self.options.is_empty() { let options = self.options.kv_pairs(); write!(f, "WITH(\n{}\n)", format_list_indent!(options))?; @@ -688,8 +636,8 @@ impl Display for CreateFlow { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - writeln!(f, "{}", &self.flow_name)?; - writeln!(f, "SINK TO {}", &self.sink_table_name)?; + writeln!(f, "{}", self.flow_name)?; + writeln!(f, "SINK TO {}", self.sink_table_name)?; if let Some(expire_after) = &self.expire_after { writeln!(f, "EXPIRE AFTER '{} s'", expire_after)?; } @@ -710,7 +658,7 @@ impl Display for CreateFlow { let options = self.flow_options.kv_pairs(); writeln!(f, "WITH ({})", format_list_comma!(options))?; } - write!(f, "AS {}", &self.query) + write!(f, "AS {}", self.query) } } @@ -740,11 +688,11 @@ impl Display for CreateView { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - write!(f, "{} ", &self.name)?; + write!(f, "{} ", self.name)?; if !self.columns.is_empty() { write!(f, "({}) ", format_list_comma!(self.columns))?; } - write!(f, "AS {}", &self.query) + write!(f, "AS {}", self.query) } } @@ -761,6 +709,120 @@ mod tests { use crate::parser::{ParseOptions, ParserContext}; use crate::statements::statement::Statement; + #[test] + fn test_rewritten_create_sql_budget() { + let parse = |sql| { + let mut statements = ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::CreateTable(create) = statements.remove(0) else { + panic!("CREATE expected") + }; + create + }; + let create = parse( + "CREATE TABLE demo (ts TIMESTAMP TIME INDEX, host STRING) ENGINE=metric WITH (on_physical_table='phy')", + ); + let size = create.to_string().len(); + let mut remaining = size; + assert!(create.consume_sql_size_budget(&mut remaining)); + assert_eq!(remaining, 0); + assert!(!create.consume_sql_size_budget(&mut (size - 1))); + + for location in ["table", "fulltext", "skipping", "inverted"] { + for value_type in ["secret", "array", "struct"] { + for length in [10, 8192] { + let mut create = create.clone(); + let extensions = &mut create.columns[1].extensions; + let map = match location { + "table" => &mut create.options, + "fulltext" => extensions.fulltext_index_options.get_or_insert_default(), + "skipping" => extensions.skipping_index_options.get_or_insert_default(), + _ => extensions.inverted_index_options.get_or_insert_default(), + }; + let value = "private-value".repeat(length); + match value_type { + "secret" => map.insert("secret_access_key".into(), value.clone()), + "array" => map.insert_options("payload", vec![value.as_str()].into()), + _ => map.insert_options( + "payload", + crate::util::OptionValue::try_new(Expr::Struct { + values: vec![Expr::Value( + sqlparser::ast::Value::SingleQuotedString(value.clone()).into(), + )], + fields: vec![], + }) + .unwrap(), + ), + } + assert!(!create.to_string().contains(&value)); + assert_eq!( + create.consume_sql_size_budget(&mut 4096), + length == 10, + "{location}/{value_type}" + ); + } + } + } + let mut vector = parse("CREATE TABLE demo (ts TIMESTAMP TIME INDEX, v VECTOR(3))"); + vector.columns[1].column_def.options.push(ColumnOptionDef { + name: None, + option: sqlparser::ast::ColumnOption::Comment("private-value".repeat(8192)), + }); + assert!(vector.to_string().contains("private-value")); + assert!(!vector.consume_sql_size_budget(&mut 4096)); + vector.columns[1].column_def.options.clear(); + assert!(vector.consume_sql_size_budget(&mut 4096)); + if let DataType::Custom(_, tokens) = &mut vector.columns[1].column_def.data_type { + tokens[0] = format!("{}3", "0".repeat(8192)); + } + assert!(!vector.consume_sql_size_budget(&mut 4096)); + vector.columns[1] + .extensions + .vector_options + .as_mut() + .unwrap() + .insert("extra".into(), "x".into()); + assert!(!vector.consume_sql_size_budget(&mut 4096)); + + let mut create = create; + create.columns[1].extensions.json2_options = Some(Json2Options { + max_auto_expanded_paths: Some(1), + type_hints: vec![JsonTypeHint { + path: vec!["x".repeat(8192)], + data_type: DataType::Text, + inverted_index: true, + }], + }); + assert!(!create.consume_sql_size_budget(&mut 4096)); + create.columns[1].extensions.json2_options = None; + create.partitions = Some(Partitions { + column_list: vec![], + exprs: vec![Expr::Identifier(Ident::new("hidden"))], + }); + assert!(!create.consume_sql_size_budget(&mut 4096)); + } + + #[test] + fn test_json2_streamed_quote_escaping() { + let path = "a\\\"\"b"; + let options = Json2Options { + max_auto_expanded_paths: None, + type_hints: vec![JsonTypeHint { + path: vec![path.into()], + data_type: DataType::Text, + inverted_index: false, + }], + }; + assert_eq!( + options.to_string(), + format!("(\n \"{}\" TEXT\n )", path.replace('"', "\"\"")) + ); + } + #[test] fn test_display_create_table() { let sql = r"create table if not exists demo( @@ -930,7 +992,7 @@ ENGINE=mito let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( "service.name" STRING, - "a.b"."c" INT64, + "a.b"."c" BIGINT, a."b.c" STRING ), ts TIMESTAMP TIME INDEX @@ -974,7 +1036,7 @@ ENGINE=mito fn test_parse_json2_max_auto_expanded_paths_option() -> Result<()> { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( - status_code INT64, + status_code BIGINT, max_auto_expanded_paths = 1 ), ts TIMESTAMP TIME INDEX @@ -1000,7 +1062,7 @@ ENGINE=mito let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( "1abc" STRING, - a."2b" INT64 + a."2b" BIGINT ), ts TIMESTAMP TIME INDEX )"#; @@ -1042,8 +1104,8 @@ ENGINE=mito fn test_json2_type_hint_rejects_default() { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( - status_code INT64 DEFAULT -5, - duration FLOAT64 DEFAULT +1.5, + status_code BIGINT DEFAULT -5, + duration DOUBLE DEFAULT +1.5, error BOOLEAN DEFAULT false, message STRING DEFAULT 'unknown' ), @@ -1059,7 +1121,7 @@ ENGINE=mito fn test_json2_type_hint_rejects_not_null() { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( - status_code INT64 NOT NULL DEFAULT NULL + status_code BIGINT NOT NULL DEFAULT NULL ), ts TIMESTAMP TIME INDEX )"#; @@ -1070,23 +1132,23 @@ ENGINE=mito } #[test] - fn test_set_json_settings_normalizes_type_hint_sql_types() -> Result<()> { + fn test_set_json_settings_preserves_type_hint_sql_types() -> Result<()> { let mut extensions = super::ColumnExtensions::default(); let settings = JsonSettings::try_new( vec![ DatatypeJsonTypeHint { path: vec!["i".to_string()], - data_type: ConcreteDataType::int32_datatype(), + data_type: ConcreteDataType::int64_datatype(), inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["f".to_string()], - data_type: ConcreteDataType::float32_datatype(), + data_type: ConcreteDataType::float64_datatype(), inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["u".to_string()], - data_type: ConcreteDataType::uint32_datatype(), + data_type: ConcreteDataType::uint64_datatype(), inverted_index: false, }, DatatypeJsonTypeHint { @@ -1312,76 +1374,4 @@ AS SELECT number FROM numbers_input where number > 10"#, _ => unreachable!(), } } - - #[test] - fn test_vector_index_options_validation() { - use super::{ColumnExtensions, OptionMap}; - - // Test zero connectivity should fail - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([( - "connectivity".to_string(), - "0".to_string(), - )])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("connectivity must be in the range [2, 2048]") - ); - - // Test zero expansion_add should fail - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([( - "expansion_add".to_string(), - "0".to_string(), - )])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("expansion_add must be greater than 0") - ); - - // Test zero expansion_search should fail - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([( - "expansion_search".to_string(), - "0".to_string(), - )])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("expansion_search must be greater than 0") - ); - - // Test valid values should succeed - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([ - ("connectivity".to_string(), "32".to_string()), - ("expansion_add".to_string(), "200".to_string()), - ("expansion_search".to_string(), "100".to_string()), - ])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_ok()); - let options = result.unwrap().unwrap(); - assert_eq!(options.connectivity, 32); - assert_eq!(options.expansion_add, 200); - assert_eq!(options.expansion_search, 100); - } } diff --git a/src/sql/src/statements/insert.rs b/src/sql/src/statements/insert.rs index 52318cddbf4..9c069bde43d 100644 --- a/src/sql/src/statements/insert.rs +++ b/src/sql/src/statements/insert.rs @@ -14,7 +14,7 @@ use serde::Serialize; use sqlparser::ast::{ - Insert as SpInsert, ObjectName, ObjectNamePart, Parens, Query, SetExpr, Statement, TableObject, + Insert as SpInsert, ObjectName, ObjectNamePart, Parens, SetExpr, Statement, TableObject, UnaryOperator, ValueWithSpan, Values, }; use sqlparser::parser::ParserError; @@ -70,14 +70,18 @@ impl Insert { /// Extracts the literal insert statement body if possible pub fn values_body(&self) -> Result>> { match &self.inner { - Statement::Insert(SpInsert { - source: - Some(box Query { - body: box SetExpr::Values(Values { rows, .. }), - .. - }), - .. - }) => sql_exprs_to_values(rows), + Statement::Insert(SpInsert { source, .. }) => { + let rows = source + .as_deref() + .and_then(|query| match query.body.as_ref() { + SetExpr::Values(Values { rows, .. }) => Some(rows), + _ => None, + }); + match rows { + Some(rows) => sql_exprs_to_values(rows), + None => unreachable!(), + } + } _ => unreachable!(), } } @@ -86,36 +90,37 @@ impl Insert { /// The rules is the same as function `values_body()`. pub fn can_extract_values(&self) -> bool { match &self.inner { - Statement::Insert(SpInsert { - source: - Some(box Query { - body: box SetExpr::Values(Values { rows, .. }), - .. - }), - .. - }) => rows.iter().all(|es| { - es.iter().all(|expr| match expr { - Expr::Value(_) => true, - Expr::Identifier(ident) => { - if ident.quote_style.is_none() { - ident.value.to_lowercase() == "default" - } else { - ident.quote_style == Some('"') - } - } - Expr::UnaryOp { op, expr } => { - matches!(op, UnaryOperator::Minus | UnaryOperator::Plus) - && matches!( - &**expr, - Expr::Value(ValueWithSpan { - value: Value::Number(_, _), - .. - }) - ) - } - _ => false, + Statement::Insert(SpInsert { source, .. }) => source + .as_deref() + .and_then(|query| match query.body.as_ref() { + SetExpr::Values(Values { rows, .. }) => Some(rows), + _ => None, }) - }), + .is_some_and(|rows| { + rows.iter().all(|es| { + es.iter().all(|expr| match expr { + Expr::Value(_) => true, + Expr::Identifier(ident) => { + if ident.quote_style.is_none() { + ident.value.to_lowercase() == "default" + } else { + ident.quote_style == Some('"') + } + } + Expr::UnaryOp { op, expr } => { + matches!(op, UnaryOperator::Minus | UnaryOperator::Plus) + && matches!( + &**expr, + Expr::Value(ValueWithSpan { + value: Value::Number(_, _), + .. + }) + ) + } + _ => false, + }) + }) + }), _ => false, } } @@ -124,7 +129,7 @@ impl Insert { pub fn has_non_values_query_source(&self) -> bool { match &self.inner { Statement::Insert(SpInsert { - source: Some(box query), + source: Some(query), .. }) => !matches!(&*query.body, SetExpr::Values(_)), _ => false, @@ -134,9 +139,9 @@ impl Insert { pub fn query_body(&self) -> Result> { Ok(match &self.inner { Statement::Insert(SpInsert { - source: Some(box query), + source: Some(query), .. - }) => Some(query.clone().try_into()?), + }) => Some(query.as_ref().clone().try_into()?), _ => None, }) } @@ -383,11 +388,9 @@ mod tests { let q = insert.query_body().unwrap().unwrap(); assert!(insert.has_non_values_query_source()); assert!(matches!( - q.inner, - Query { - body: box SetExpr::Select { .. }, - .. - } + &q.inner, + sqlparser::ast::Query { body, .. } + if matches!(body.as_ref(), SetExpr::Select { .. }) )); } _ => unreachable!(), diff --git a/src/sql/src/statements/option_map.rs b/src/sql/src/statements/option_map.rs index d562a61ce04..29808ae602d 100644 --- a/src/sql/src/statements/option_map.rs +++ b/src/sql/src/statements/option_map.rs @@ -139,26 +139,44 @@ impl OptionMap { } pub fn kv_pairs(&self) -> Vec { - let mut result = Vec::with_capacity(self.options.len() + self.secrets.len()); - for (k, v) in self - .options - .iter() - .filter_map(|(k, v)| v.as_string().map(|v| (k, v))) - { - if k.contains(".") { - result.push(format!("'{k}' = '{}'", v.escape_debug())); - } else { - result.push(format!("{k} = '{}'", v.escape_debug())); + self.sql_pairs(false).map(|pair| pair.to_string()).collect() + } + + /// Streams the same pairs as `kv_pairs`, or all actual values for size validation. + pub(crate) fn write_sql( + &self, + f: &mut dyn std::fmt::Write, + separator: &str, + prefix: &str, + complete: bool, + ) -> std::fmt::Result { + for (index, pair) in self.sql_pairs(complete).enumerate() { + if index > 0 { + f.write_str(separator)?; } + write!(f, "{prefix}{pair}")?; } - for (k, _) in self.secrets.iter() { - if k.contains(".") { - result.push(format!("'{k}' = '******'")); + Ok(()) + } + + fn sql_pairs(&self, complete: bool) -> impl Iterator> { + let options = self.options.iter().filter_map(move |(key, value)| { + let value = match value.as_string() { + Some(value) => Either::Right(value), + None if complete => Either::Left(value), + None => return None, + }; + Some(SqlPair { key, value }) + }); + let secrets = self.secrets.iter().map(move |(key, value)| SqlPair { + key, + value: Either::Right(if complete { + value.expose_secret() } else { - result.push(format!("{k} = '******'")); - } - } - result + "******" + }), + }); + options.chain(secrets) } pub fn entries(&self) -> impl Iterator)> { @@ -174,6 +192,25 @@ impl OptionMap { } } +struct SqlPair<'a> { + key: &'a str, + value: Either<&'a OptionValue, &'a str>, +} + +impl std::fmt::Display for SqlPair<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.key.contains('.') { + write!(f, "'{}'", self.key)?; + } else { + f.write_str(self.key)?; + } + match self.value { + Either::Left(value) => write!(f, " = {value}"), + Either::Right(value) => write!(f, " = '{}'", value.escape_debug()), + } + } +} + impl> From for OptionMap { fn from(value: I) -> Self { let mut result = OptionMap::default(); @@ -221,10 +258,10 @@ impl Visit for OptionMap { impl VisitMut for OptionMap { fn visit(&mut self, visitor: &mut V) -> ControlFlow { - for (_, v) in self.options.iter_mut() { + for v in self.options.values_mut() { v.visit(visitor)?; } - for (_, v) in self.secrets.iter_mut() { + for v in self.secrets.values_mut() { v.expose_secret_mut().visit(visitor)?; } ControlFlow::Continue(()) diff --git a/src/sql/src/statements/show.rs b/src/sql/src/statements/show.rs index 67624192b82..d1531de360d 100644 --- a/src/sql/src/statements/show.rs +++ b/src/sql/src/statements/show.rs @@ -41,7 +41,7 @@ impl Display for ShowKind { macro_rules! format_kind { ($self: expr, $f: expr) => { if $self.kind != ShowKind::All { - write!($f, " {}", &$self.kind)?; + write!($f, " {}", $self.kind)?; } }; } @@ -71,7 +71,7 @@ impl Display for ShowColumns { if self.full { write!(f, "FULL ")?; } - write!(f, "COLUMNS IN {}", &self.table)?; + write!(f, "COLUMNS IN {}", self.table)?; if let Some(database) = &self.database { write!(f, " IN {database}")?; } @@ -90,7 +90,7 @@ pub struct ShowIndex { impl Display for ShowIndex { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "SHOW INDEX IN {}", &self.table)?; + write!(f, "SHOW INDEX IN {}", self.table)?; if let Some(database) = &self.database { write!(f, " IN {database}")?; } @@ -110,7 +110,7 @@ pub struct ShowRegion { impl Display for ShowRegion { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "SHOW REGION IN {}", &self.table)?; + write!(f, "SHOW REGION IN {}", self.table)?; if let Some(database) = &self.database { write!(f, " IN {database}")?; } diff --git a/src/sql/src/statements/transform/type_alias.rs b/src/sql/src/statements/transform/type_alias.rs index b2d1803fb37..90157901f3c 100644 --- a/src/sql/src/statements/transform/type_alias.rs +++ b/src/sql/src/statements/transform/type_alias.rs @@ -72,7 +72,7 @@ impl TransformRule for TypeAliasTransformRule { alter_table.alter_operation_mut() { for add_column in add_columns { - replace_type_alias(&mut add_column.column_def.data_type); + replace_type_alias(add_column.column.mut_data_type()); } } } diff --git a/src/sql/src/util.rs b/src/sql/src/util.rs index c0c237cb0c0..a33c2576bbe 100644 --- a/src/sql/src/util.rs +++ b/src/sql/src/util.rs @@ -16,7 +16,6 @@ use std::collections::HashSet; use std::fmt::{Display, Formatter}; use std::ops::ControlFlow; -use itertools::Itertools; use promql_parser::label::{METRIC_NAME, MatchOp}; use promql_parser::parser::{ Expr as PromExpr, MatrixSelector as PromMatrixSelector, VectorSelector as PromVectorSelector, @@ -161,12 +160,20 @@ impl Display for OptionValue { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { if let Some(s) = self.as_string() { write!(f, "'{s}'") - } else if let Some(s) = self.as_list() { - write!( - f, - "[{}]", - s.into_iter().map(|x| format!("'{x}'")).join(", ") - ) + } else if let Expr::Array(array) = &self.0 + && array + .elem + .iter() + .all(|expr| Self::expr_as_string(expr).is_some()) + { + f.write_str("[")?; + for (index, expr) in array.elem.iter().enumerate() { + if index > 0 { + f.write_str(", ")?; + } + write!(f, "'{}'", Self::expr_as_string(expr).unwrap())?; + } + f.write_str("]") } else { write!(f, "'{}'", self.0) } @@ -667,6 +674,7 @@ fn extract_tables_from_sql_query(query: &sqlparser::ast::Query, names: &mut Hash #[cfg(test)] mod tests { + use itertools::Itertools; use sqlparser::tokenizer::Token; use super::*; diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 80ba4cf11a5..fbfcc2ccc1d 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -58,7 +58,7 @@ pub struct StandaloneOptions { pub postgres: PostgresOptions, pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, - /// Shared experimental ordinary-table batching; independent of Prom batching. + /// Ordinary-table batching with independent logical-table controls. pub pending_rows_batcher: PendingRowsBatcherOptions, pub jaeger: JaegerOptions, pub otlp: OtlpOptions, @@ -137,6 +137,7 @@ impl Configurable for StandaloneOptions { "wal.broker_endpoints", "event_recorder.event_types", "pending_rows_batcher.protocols", + "pending_rows_batcher.logical_table.protocols", ]) } } @@ -216,9 +217,62 @@ mod tests { use std::sync::Arc; use common_event_recorder::EventTypeFilter; + use servers::batcher::BatchingProtocol; use crate::options::*; + #[test] + fn test_logical_batcher_config_forwarding() { + let opts: StandaloneOptions = toml::from_str("[pending_rows_batcher]\nprotocols = ['influxdb']\n[pending_rows_batcher.logical_table]\nprotocols = ['otlp', 'prom']\npending_rows_flush_interval = '10ms'").unwrap(); + let frontend = opts.frontend_options(); + assert_eq!(frontend.pending_rows_batcher, opts.pending_rows_batcher); + assert_eq!( + frontend.pending_rows_batcher.logical_table, + opts.pending_rows_batcher.logical_table + ); + let restored: StandaloneOptions = toml::from_str(&toml::to_string(&opts).unwrap()).unwrap(); + assert_eq!( + restored.pending_rows_batcher.logical_table, + opts.pending_rows_batcher.logical_table + ); + assert!( + toml::from_str::( + "[pending_rows_batcher.logical_table]\nprotocols = ['logs']" + ) + .is_err() + ); + } + + #[test] + fn test_logical_batcher_protocols_from_env() { + temp_env::with_vars( + [ + ( + "STANDALONE_LOGICAL_TEST__PENDING_ROWS_BATCHER__PROTOCOLS", + Some("otlp,influxdb"), + ), + ( + "STANDALONE_LOGICAL_TEST__PENDING_ROWS_BATCHER__LOGICAL_TABLE__PROTOCOLS", + Some("prom,otlp"), + ), + ], + || { + let options = + StandaloneOptions::load_layered_options(None, "STANDALONE_LOGICAL_TEST") + .unwrap(); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); + assert_eq!( + options + .pending_rows_batcher + .logical_table + .unwrap() + .protocols, + vec![BatchingProtocol::Prom, BatchingProtocol::Otlp] + ); + }, + ); + } + #[test] fn test_batcher_protocols_from_env() { temp_env::with_vars( @@ -231,11 +285,8 @@ mod tests { StandaloneOptions::load_layered_options(None, "STANDALONE_BATCHER_TEST") .unwrap(); assert_eq!( - options.pending_rows_batcher.protocols, - vec![ - servers::http::BatchingProtocol::Influxdb, - servers::http::BatchingProtocol::HttpSql - ] + options.pending_rows_batcher.table.protocols, + vec![BatchingProtocol::Influxdb, BatchingProtocol::HttpSql] ); }, ); @@ -247,6 +298,7 @@ mod tests { assert!( !defaults .pending_rows_batcher + .table .pending_rows_batching_enabled() ); let options: StandaloneOptions = toml::from_str( @@ -259,16 +311,22 @@ flow_notification_queue_capacity = 17 "#, ) .unwrap(); - assert_eq!(options.pending_rows_batcher.max_batch_rows, 25); - assert_eq!(options.pending_rows_batcher.protocols.len(), 2); + assert_eq!(options.pending_rows_batcher.table.max_batch_rows, 25); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); assert_eq!( options .pending_rows_batcher + .table .flow_notification_queue_capacity .get(), 17 ); - assert!(options.pending_rows_batcher.pending_rows_batching_enabled()); + assert!( + options + .pending_rows_batcher + .table + .pending_rows_batching_enabled() + ); let serialized = toml::to_string(&options).unwrap(); let parsed: StandaloneOptions = toml::from_str(&serialized).unwrap(); assert_eq!(options.influxdb, parsed.influxdb); diff --git a/src/store-api/Cargo.toml b/src/store-api/Cargo.toml index 61ce914fe13..4a3616afb8a 100644 --- a/src/store-api/Cargo.toml +++ b/src/store-api/Cargo.toml @@ -10,6 +10,7 @@ workspace = true [dependencies] api.workspace = true aquamarine.workspace = true +arrow-schema.workspace = true async-trait.workspace = true bytes.workspace = true common-base.workspace = true @@ -35,15 +36,10 @@ prost.workspace = true serde.workspace = true serde_json.workspace = true snafu.workspace = true -sqlparser.workspace = true strum.workspace = true tokio.workspace = true uuid.workspace = true [dev-dependencies] -async-stream.workspace = true common-meta.workspace = true serde_json.workspace = true - -[package.metadata.cargo-udeps.ignore] -normal = ["aquamarine"] diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index a2c1e3ee803..6084f207641 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -105,6 +105,24 @@ pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { /// Returns the latest entry id in the log store. fn latest_entry_id(&self, provider: &Provider) -> Result; + + /// Waits until every entry of the provider's region with an id at or below + /// `entry_id` has the persistence this log store guarantees for an + /// acknowledged append, so that a caller can persist `entry_id` as a + /// replay watermark. + /// + /// The default returns at once and keeps each log store's own + /// acknowledgement guarantee, whatever it is: Raft Engine with + /// `sync_write = false`, for example, acknowledges an append before its + /// periodic sync. A log store that acknowledges an append before its + /// entries have that persistence overrides this method. + async fn wait_durable( + &self, + _provider: &Provider, + _entry_id: EntryId, + ) -> Result<(), Self::Error> { + Ok(()) + } } /// The response of an `append` operation. diff --git a/src/store-api/src/manifest/action.rs b/src/store-api/src/manifest/action.rs deleted file mode 100644 index c9b0d3e7d0c..00000000000 --- a/src/store-api/src/manifest/action.rs +++ /dev/null @@ -1,114 +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. - -//! Common actions for manifest -use serde::{Deserialize, Serialize}; - -use crate::manifest::ManifestVersion; - -pub type ProtocolVersion = u16; - -/// Current reader and writer versions -/// TODO(dennis): configurable -const READER_VERSION: ProtocolVersion = 0; -const WRITER_VERSION: ProtocolVersion = 0; - -/// The maximum protocol version we are currently allowed to use, -/// TODO(dennis): reading from configuration. -pub fn supported_protocol_version() -> (ProtocolVersion, ProtocolVersion) { - (READER_VERSION, WRITER_VERSION) -} - -/// Protocol action that used to block older clients from reading or writing the log when backwards -/// incompatible changes are made to the protocol. -/// -/// clients should be tolerant of messages and fields that they do not understand. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct ProtocolAction { - pub min_reader_version: ProtocolVersion, - pub min_writer_version: ProtocolVersion, -} - -impl std::fmt::Display for ProtocolAction { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "Protocol({}, {})", - &self.min_reader_version, &self.min_writer_version, - ) - } -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -pub struct VersionHeader { - pub prev_version: ManifestVersion, -} - -impl Default for ProtocolAction { - fn default() -> Self { - let (min_reader_version, min_writer_version) = supported_protocol_version(); - Self { - min_reader_version, - min_writer_version, - } - } -} - -impl ProtocolAction { - pub fn new() -> Self { - Self::default() - } - - pub fn is_readable(&self, reader_version: ProtocolVersion) -> bool { - reader_version >= self.min_reader_version - } - - pub fn is_writable(&self, writer_version: ProtocolVersion) -> bool { - writer_version >= self.min_writer_version - } -} - -#[cfg(test)] -mod tests { - use serde_json as json; - - use super::*; - - #[test] - fn test_protocol_action() { - let mut action = ProtocolAction::new(); - - assert!(action.is_readable(0)); - assert!(action.is_writable(0)); - action.min_reader_version = 2; - action.min_writer_version = 3; - assert!(!action.is_readable(0)); - assert!(!action.is_writable(0)); - assert!(action.is_readable(2)); - assert!(action.is_writable(3)); - assert!(action.is_readable(3)); - assert!(action.is_writable(4)); - - let s = json::to_string(&action).unwrap(); - assert_eq!("{\"min_reader_version\":2,\"min_writer_version\":3}", s); - - let action_decoded: ProtocolAction = json::from_str(&s).unwrap(); - assert!(!action_decoded.is_readable(0)); - assert!(!action_decoded.is_writable(0)); - assert!(action_decoded.is_readable(2)); - assert!(action_decoded.is_writable(3)); - assert!(action_decoded.is_readable(3)); - assert!(action_decoded.is_writable(4)); - } -} diff --git a/src/store-api/src/manifest/storage.rs b/src/store-api/src/manifest/storage.rs deleted file mode 100644 index 58a3d27fe83..00000000000 --- a/src/store-api/src/manifest/storage.rs +++ /dev/null @@ -1,78 +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 async_trait::async_trait; -use common_error::ext::ErrorExt; - -use crate::manifest::ManifestVersion; - -#[async_trait] -pub trait LogIterator: Send + Sync { - type Error: ErrorExt + Send + Sync; - - async fn next_log(&mut self) -> Result)>, Self::Error>; -} - -#[async_trait] -pub trait ManifestLogStorage { - type Error: ErrorExt + Send + Sync; - type Iter: LogIterator; - - /// Scan the logs in [start, end) - async fn scan( - &self, - start: ManifestVersion, - end: ManifestVersion, - ) -> Result; - - /// Delete logs and checkpoints which version is less than specified `end`. - /// Keep the last checkpoint and remaining logs if `keep_last_checkpoint` is true. - /// - /// Returns the delete logs number. - async fn delete_until( - &self, - end: ManifestVersion, - keep_last_checkpoint: bool, - ) -> Result; - - /// Delete all logs and checkpoints, and remove the manifest directory. - /// The delta file corresponding to the `remove_action_version` will be deleted along with the manifest directory at the end. - async fn delete_all(&self, remove_action_version: ManifestVersion) -> Result<(), Self::Error>; - - /// Save a log - async fn save(&self, version: ManifestVersion, bytes: &[u8]) -> Result<(), Self::Error>; - - /// Delete logs in [start, end) and ignore checkpoints. - async fn delete(&self, start: ManifestVersion, end: ManifestVersion) - -> Result<(), Self::Error>; - - /// Save a checkpoint. - async fn save_checkpoint( - &self, - version: ManifestVersion, - bytes: &[u8], - ) -> Result<(), Self::Error>; - - /// Load the latest checkpoint - async fn load_last_checkpoint(&self) - -> Result)>, Self::Error>; - /// Delete the checkpoint by version - async fn delete_checkpoint(&self, version: ManifestVersion) -> Result<(), Self::Error>; - - /// Load the checkpoint by version - async fn load_checkpoint( - &self, - version: ManifestVersion, - ) -> Result)>, Self::Error>; -} diff --git a/src/store-api/src/metadata.rs b/src/store-api/src/metadata.rs index 0c663bccc05..6983bafb504 100644 --- a/src/store-api/src/metadata.rs +++ b/src/store-api/src/metadata.rs @@ -30,7 +30,8 @@ use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; use datatypes::arrow; use datatypes::arrow::datatypes::FieldRef; -use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef, VectorIndexOptions}; +use datatypes::extension::json::json2_metadata_with_updated_settings; +use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef}; use datatypes::types::TimestampType; use itertools::Itertools; use serde::de::Error; @@ -422,22 +423,6 @@ impl RegionMetadata { inverted_index } - /// Gets the column IDs that have vector indexes along with their options. - /// Returns a map from column ID to the vector index options. - pub fn vector_indexed_column_ids(&self) -> HashMap { - self.column_metadatas - .iter() - .filter_map(|column| { - column - .column_schema - .vector_index_options() - .ok() - .flatten() - .map(|options| (column.column_id, options)) - }) - .collect() - } - /// Checks whether the metadata is valid. fn validate(&self) -> Result<()> { // Id to name. @@ -663,6 +648,10 @@ impl RegionMetadataBuilder { AlterKind::AddColumns { columns } => self.add_columns(columns)?, AlterKind::DropColumns { names } => self.drop_columns(&names), AlterKind::ModifyColumnTypes { columns } => self.modify_column_types(columns)?, + AlterKind::SetJsonSettings { + column_name, + settings, + } => self.set_json_settings(column_name, settings)?, AlterKind::SetIndexes { options } => self.set_indexes(options)?, AlterKind::UnsetIndexes { options } => self.unset_indexes(options)?, AlterKind::SetRegionOptions { options: _ } => { @@ -831,6 +820,37 @@ impl RegionMetadataBuilder { Ok(()) } + fn set_json_settings( + &mut self, + col_name: String, + settings: datatypes::json::JsonSettings, + ) -> Result<()> { + let Some(col_meta) = self + .column_metadatas + .iter_mut() + .find(|col| col.column_schema.name == col_name) + else { + return InvalidRegionRequestSnafu { + region_id: self.region_id, + err: format!("column {col_name} not found"), + } + .fail(); + }; + + let old_metadata = col_meta.column_schema.metadata(); + let new_metadata = + json2_metadata_with_updated_settings(old_metadata, settings).map_err(|err| { + InvalidRegionRequestSnafu { + region_id: self.region_id, + err: err.to_string(), + } + .build() + })?; + + *col_meta.column_schema.mut_metadata() = new_metadata; + Ok(()) + } + fn set_indexes(&mut self, options: Vec) -> Result<()> { let mut set_index_map: HashMap<_, Vec<_>> = HashMap::new(); for option in &options { diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index bc9b31ada95..07cfe6a6516 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -29,8 +29,6 @@ pub const MERGE_MODE_KEY: &str = "merge_mode"; pub const TTL_KEY: &str = "ttl"; /// Option key for the per-table auto flush interval. pub const AUTO_FLUSH_INTERVAL_KEY: &str = "auto_flush_interval"; -/// Option key for snapshot read. -pub const SNAPSHOT_READ: &str = "snapshot_read"; /// Option key for compaction type. pub const COMPACTION_TYPE: &str = "compaction.type"; /// Option key for forcing compaction options override. diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 826ef7d6bc9..e1f0eb22497 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -26,12 +26,14 @@ use api::v1::region::{ AlterRequest, AlterRequests, BuildIndexRequest, BulkInsertRequest, CleanUpRequest as PbCleanUpRequest, CloseRequest, CompactRequest, CreateRequest, CreateRequests, DeleteRequests, DropRequest, DropRequests, FlushRequest, InsertRequests, - OpenRequest, TruncateRequest, alter_request, compact_request, region_request, truncate_request, + OpenRequest, TruncateRequest, alter_request, build_index_request, compact_request, + region_request, truncate_request, }; use api::v1::{ self, Analyzer, ArrowIpc, FulltextBackend as PbFulltextBackend, Option as PbOption, Rows, SemanticType, SkippingIndexType as PbSkippingIndexType, WriteHint, }; +use arrow_schema::extension::ExtensionType; pub use common_base::AffectedRows; use common_base::readable_size::ReadableSize; use common_grpc::flight::FlightDecoder; @@ -39,6 +41,8 @@ use common_recordbatch::DfRecordBatch; use common_time::range::TimestampRange; use common_time::{TimeToLive, Timestamp}; use datatypes::error::time_index_not_widening_error; +use datatypes::extension::json::Json2ExtensionType; +use datatypes::json::{JsonSettings, JsonTypeHint}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use num_enum::TryFromPrimitive; @@ -63,7 +67,7 @@ use crate::mito_engine_options::{ TWCS_MAX_OUTPUT_FILE_SIZE, TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, WRITE_BUFFER_SIZE_KEY, }; use crate::path_utils::table_dir; -use crate::storage::{ColumnId, RegionId, ScanRequest}; +use crate::storage::{ColumnId, RegionId}; /// The type of path to generate. #[derive(Debug, Clone, Copy, PartialEq, TryFromPrimitive)] @@ -439,7 +443,9 @@ fn make_region_build_index(index: BuildIndexRequest) -> Result, } -#[derive(Debug)] -pub struct RegionReadRequest { - pub request: ScanRequest, -} - /// Request to delete data from a region. #[derive(Debug)] pub struct RegionDeleteRequest { @@ -795,6 +796,13 @@ pub enum AlterKind { /// Columns to change. columns: Vec, }, + /// Set JSON2 settings of a region column. + SetJsonSettings { + /// Column name. + column_name: String, + /// Target JSON2 settings. + settings: JsonSettings, + }, /// Set region options. SetRegionOptions { options: Vec }, /// Unset region options. @@ -975,6 +983,9 @@ impl AlterKind { col_to_change.validate(metadata)?; } } + AlterKind::SetJsonSettings { column_name, .. } => { + Self::validate_set_json_settings(column_name, metadata)? + } AlterKind::SetRegionOptions { .. } => {} AlterKind::UnsetRegionOptions { .. } => {} AlterKind::SetIndexes { options } => { @@ -1084,6 +1095,18 @@ impl AlterKind { AlterKind::ModifyColumnTypes { columns } => columns .iter() .any(|col_to_change| col_to_change.need_alter(metadata)), + AlterKind::SetJsonSettings { + column_name, + settings, + } => metadata.column_by_name(column_name).is_some_and(|col| { + col.column_schema + .extension_type::() + .ok() + .flatten() + .is_none_or(|extension| { + !extension.metadata().json_settings().equivalent(settings) + }) + }), AlterKind::SetRegionOptions { .. } => true, AlterKind::UnsetRegionOptions { .. } => true, AlterKind::SetIndexes { options, .. } => options @@ -1160,6 +1183,34 @@ impl AlterKind { Ok(()) } + + fn validate_set_json_settings(col_name: &String, metadata: &RegionMetadata) -> Result<()> { + let region_id = metadata.region_id; + + let col = metadata + .column_by_name(col_name) + .with_context(|| InvalidRegionRequestSnafu { + region_id, + err: format!("column {} not found", col_name), + })?; + + ensure!( + col.semantic_type == SemanticType::Field, + InvalidRegionRequestSnafu { + region_id, + err: format!("column {} is not a field column", col_name), + } + ); + ensure!( + col.column_schema.data_type.is_json2(), + InvalidRegionRequestSnafu { + region_id, + err: format!("column {} is not a JSON2 column", col_name), + } + ); + + Ok(()) + } } impl TryFrom for AlterKind { @@ -1183,6 +1234,15 @@ impl TryFrom for AlterKind { .collect::>(); AlterKind::ModifyColumnTypes { columns } } + alter_request::Kind::SetJsonSettings(x) => { + let settings = x.settings.context(InvalidRawRegionRequestSnafu { + err: "missing settings in SetJsonSettings", + })?; + AlterKind::SetJsonSettings { + column_name: x.column_name, + settings: json_settings_from_proto(settings)?, + } + } alter_request::Kind::DropColumns(x) => { let names = x.drop_columns.into_iter().map(|x| x.name).collect(); AlterKind::DropColumns { names } @@ -1484,6 +1544,36 @@ impl From for ModifyColumnType { } } +fn json_settings_from_proto(settings: v1::JsonSettings) -> Result { + let type_hints = settings + .type_hints + .into_iter() + .map(|hint| { + let wrapper = ColumnDataTypeWrapper::try_new(hint.data_type, hint.datatype_extension) + .map_err(|err| { + InvalidRawRegionRequestSnafu { + err: err.to_string(), + } + .build() + })?; + let data_type = ConcreteDataType::from(wrapper); + + Ok(JsonTypeHint { + path: hint.path, + data_type, + inverted_index: false, + }) + }) + .collect::>>()?; + + JsonSettings::try_new(type_hints, settings.max_auto_expanded_paths).map_err(|err| { + InvalidRawRegionRequestSnafu { + err: err.to_string(), + } + .build() + }) +} + /// Region option changes used by ALTER requests. /// /// This type is serialized for request persistence. Keep future changes backward @@ -1796,7 +1886,10 @@ impl Default for RegionCompactRequest { } #[derive(Debug, Clone, Default)] -pub struct RegionBuildIndexRequest {} +pub struct RegionBuildIndexRequest { + /// The index build mode. Absent options select SST indexes. + pub options: Option, +} /// Truncate region request. #[derive(Debug)] @@ -1953,12 +2046,42 @@ mod tests { use api::v1::region::RegionColumnDef; use api::v1::{ColumnDataType, ColumnDef}; use common_time::range::TimestampRange; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend}; + use datatypes::types::JsonType; use super::*; use crate::metadata::RegionMetadataBuilder; + #[test] + fn test_build_index_options_round_trip() { + use prost::Message; + + for options in [ + None, + Some(build_index_request::Options::SstIndex(Default::default())), + Some(build_index_request::Options::SeriesIndex(Default::default())), + ] { + let request = BuildIndexRequest { + region_id: 42, + options, + }; + let decoded = BuildIndexRequest::decode(request.encode_to_vec().as_slice()).unwrap(); + let requests = make_region_build_index(decoded).unwrap(); + assert_eq!(1, requests.len()); + assert_eq!(RegionId::from_u64(42), requests[0].0); + let RegionRequest::BuildIndex(request) = &requests[0].1 else { + panic!("expected build-index request"); + }; + assert_eq!(options, request.options); + } + // The legacy wire message contains only region_id (field 1). + let legacy = BuildIndexRequest::decode(&[0x08, 42][..]).unwrap(); + assert!(legacy.options.is_none()); + } + #[test] fn test_make_region_puts_preserves_skip_wal() { let region_id = RegionId::new(42, 3); @@ -2364,6 +2487,34 @@ mod tests { }, } ); + + let request = RegionAlterRequest::try_from(AlterRequest { + region_id: 0, + schema_version: 1, + kind: Some(alter_request::Kind::SetJsonSettings(v1::SetJsonSettings { + column_name: "payload".to_string(), + settings: Some(v1::JsonSettings { + type_hints: vec![v1::JsonTypeHint { + path: vec!["service".to_string()], + data_type: ColumnDataType::String as i32, + datatype_extension: None, + }], + max_auto_expanded_paths: Some(10), + }), + })), + }) + .unwrap(); + + let AlterKind::SetJsonSettings { + column_name, + settings, + } = request.kind + else { + unreachable!() + }; + assert_eq!("payload", column_name); + assert_eq!(Some(10), settings.max_auto_expanded_paths()); + assert_eq!(1, settings.type_hints().len()); } #[test] @@ -2446,6 +2597,15 @@ mod tests { builder.build().unwrap() } + fn json2_column_schema(name: &str, settings: JsonSettings) -> ColumnSchema { + let mut column_schema = + ColumnSchema::new(name, ConcreteDataType::Json(JsonType::null()), true); + column_schema.with_extension_type(&Json2ExtensionType::new(std::sync::Arc::new( + JsonMetadata::new(settings), + ))); + column_schema + } + #[test] fn test_add_column_validate() { let metadata = new_metadata(); @@ -2730,6 +2890,40 @@ mod tests { assert!(kind.need_alter(&metadata)); } + #[test] + fn test_validate_set_json_settings() { + let mut metadata = new_metadata(); + let current_schema = json2_column_schema("field_0", JsonSettings::new_v2()); + let target_settings = JsonSettings::try_new(vec![], Some(10)).unwrap(); + metadata + .column_metadatas + .iter_mut() + .find(|column| column.column_schema.name == "field_0") + .unwrap() + .column_schema = current_schema.clone(); + + let kind = AlterKind::SetJsonSettings { + column_name: "field_0".to_string(), + settings: target_settings.clone(), + }; + kind.validate(&metadata).unwrap(); + assert!(kind.need_alter(&metadata)); + + let no_op = AlterKind::SetJsonSettings { + column_name: "field_0".to_string(), + settings: JsonSettings::new_v2(), + }; + no_op.validate(&metadata).unwrap(); + assert!(!no_op.need_alter(&metadata)); + + AlterKind::SetJsonSettings { + column_name: "tag_0".to_string(), + settings: target_settings, + } + .validate(&metadata) + .unwrap_err(); + } + #[test] fn test_validate_add_columns() { let kind = AlterKind::AddColumns { diff --git a/src/store-api/src/storage.rs b/src/store-api/src/storage.rs index cae2439168a..e0761f8a006 100644 --- a/src/store-api/src/storage.rs +++ b/src/store-api/src/storage.rs @@ -29,8 +29,5 @@ pub use datatypes::schema::{ pub use self::descriptors::*; pub use self::file::{FileId, FileRef, FileRefsManifest, GcReport, IndexVersion, ParseIdError}; pub use self::projection::NestedPath; -pub use self::requests::{ - ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector, VectorDistanceMetric, - VectorIndexEngine, VectorIndexEngineType, VectorSearchMatches, VectorSearchRequest, -}; +pub use self::requests::{ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; pub use self::types::{SequenceNumber, SequenceRange}; diff --git a/src/store-api/src/storage/requests.rs b/src/store-api/src/storage/requests.rs index b35a1907d43..3c9d4951d4b 100644 --- a/src/store-api/src/storage/requests.rs +++ b/src/store-api/src/storage/requests.rs @@ -15,68 +15,13 @@ use std::collections::HashMap; use std::fmt::{Display, Formatter}; -use common_error::ext::BoxedError; use common_recordbatch::OrderOption; use datafusion_expr::expr::Expr; -// Re-export vector types from datatypes to avoid duplication -pub use datatypes::schema::{VectorDistanceMetric, VectorIndexEngineType}; use datatypes::types::json_type::JsonNativeType; use itertools::Itertools; use strum::Display; -use crate::storage::{ColumnId, SequenceNumber}; - -/// A hint for KNN vector search. -#[derive(Debug, Clone, PartialEq)] -pub struct VectorSearchRequest { - /// Column ID of the vector column to search. - pub column_id: ColumnId, - /// The query vector to search for. - pub query_vector: Vec, - /// Number of nearest neighbors to return. - pub k: usize, - /// Distance metric to use (matches the index metric). - pub metric: VectorDistanceMetric, -} - -/// Search results from vector index. -#[derive(Debug, Clone, PartialEq)] -pub struct VectorSearchMatches { - /// Keys (row offsets in the index). - pub keys: Vec, - /// Distances from the query vector. - pub distances: Vec, -} - -/// Trait for vector index engines (HNSW implementations). -/// -/// This trait defines the interface for pluggable vector index engines. -/// Implementations (e.g., UsearchEngine) are provided by storage engines like mito2. -pub trait VectorIndexEngine: Send + Sync { - /// Adds a vector with the given key. - fn add(&mut self, key: u64, vector: &[f32]) -> Result<(), BoxedError>; - - /// Searches for k nearest neighbors. - fn search(&self, query: &[f32], k: usize) -> Result; - - /// Returns the serialized length. - fn serialized_length(&self) -> usize; - - /// Serializes the index to a buffer. - fn save_to_buffer(&self, buffer: &mut [u8]) -> Result<(), BoxedError>; - - /// Reserves capacity for vectors. - fn reserve(&mut self, capacity: usize) -> Result<(), BoxedError>; - - /// Returns current size (number of vectors). - fn size(&self) -> usize; - - /// Returns current capacity. - fn capacity(&self) -> usize; - - /// Returns memory usage in bytes. - fn memory_usage(&self) -> usize; -} +use crate::storage::SequenceNumber; /// A hint on how to select rows from a time-series. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Display)] @@ -144,9 +89,6 @@ pub struct ScanRequest { pub exact_sequence_range: bool, /// Optional hint for the distribution of time-series data. pub distribution: Option, - /// Optional hint for KNN vector search. When set, the scan should use - /// vector index to find the k nearest neighbors. - pub vector_search: Option, /// Optional hint from query-driven JSON type concretization. pub json_type_hint: HashMap, /// Whether Mito should keep string primary-key columns dictionary encoded in its output. @@ -249,16 +191,6 @@ impl Display for ScanRequest { if let Some(distribution) = &self.distribution { write!(f, "{}distribution: {}", delimiter.as_str(), distribution)?; } - if let Some(vector_search) = &self.vector_search { - write!( - f, - "{}vector_search: column_id={}, k={}, metric={}", - delimiter.as_str(), - vector_search.column_id, - vector_search.k, - vector_search.metric - )?; - } if !self.json_type_hint.is_empty() { write!( f, diff --git a/src/table/Cargo.toml b/src/table/Cargo.toml index 8ecc51a320c..357d005e612 100644 --- a/src/table/Cargo.toml +++ b/src/table/Cargo.toml @@ -34,13 +34,11 @@ derive_builder.workspace = true futures.workspace = true greptime-proto.workspace = true humantime.workspace = true -humantime-serde.workspace = true lazy_static.workspace = true once_cell.workspace = true paste.workspace = true serde.workspace = true snafu.workspace = true -sqlparser.workspace = true store-api.workspace = true tokio.workspace = true @@ -48,4 +46,3 @@ tokio.workspace = true common-test-util.workspace = true parquet = { workspace = true, features = ["async"] } serde_json.workspace = true -tokio-util = { version = "0.7", features = ["compat"] } diff --git a/src/table/src/lib.rs b/src/table/src/lib.rs index abc7dc83002..9968ee4dc97 100644 --- a/src/table/src/lib.rs +++ b/src/table/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(try_blocks)] - pub mod dist_table; pub mod error; pub mod metadata; diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 5916dd7ab49..c8941bb1a82 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -22,6 +22,7 @@ use common_query::AddColumnLocation; use datafusion_expr::TableProviderFilterPushDown; use datatypes::error::time_index_not_widening_error; pub use datatypes::error::{Error as ConvertError, Result as ConvertResult}; +use datatypes::extension::json::json2_metadata_with_updated_settings; use datatypes::schema::{ ColumnSchema, FulltextOptions, Schema, SchemaBuilder, SchemaRef, SkippingIndexOptions, }; @@ -41,9 +42,9 @@ use crate::error::{self, Result}; use crate::requests::{ AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError, 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, + SetDefaultRequest, SetIndexOption, SetJsonSettingsRequest, 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; @@ -332,6 +333,7 @@ impl TableMeta { AlterKind::ModifyColumnTypes { columns } => { self.modify_column_types(table_name, columns) } + AlterKind::SetJsonSettings { request } => self.set_json_settings(table_name, request), // No need to rebuild table meta when renaming tables. AlterKind::RenameTable { .. } => Ok(self.new_meta_builder()), AlterKind::SetTableOptions { options } => self.set_table_options(options), @@ -1141,7 +1143,7 @@ impl TableMeta { .with_context(|_| error::CastDefaultValueSnafu { reason: format!( "Failed to cast default value from {:?} to type {:?}", - default_value, &change_column.target_type + default_value, change_column.target_type ), })?, ) @@ -1187,6 +1189,77 @@ impl TableMeta { Ok(meta_builder) } + fn set_json_settings( + &self, + table_name: &str, + req: &SetJsonSettingsRequest, + ) -> Result { + let table_schema = &self.schema; + let idx = table_schema + .column_index_by_name(&req.column_name) + .with_context(|| error::ColumnNotExistsSnafu { + column_name: &req.column_name, + table_name, + })?; + let col = &table_schema.column_schemas()[idx]; + + ensure!( + !self.primary_key_indices.contains(&idx) && table_schema.timestamp_index() != Some(idx), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "Not allowed to change JSON settings for key or timestamp column '{}'", + col.name + ), + } + ); + ensure!( + col.data_type.is_json2(), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!("column '{}' is not a JSON2 column", col.name), + } + ); + + let target_metadata = + json2_metadata_with_updated_settings(col.metadata(), req.settings.clone()).map_err( + |err| { + error::InvalidAlterRequestSnafu { + table: table_name, + err: err.to_string(), + } + .build() + }, + )?; + + let mut cols = table_schema.column_schemas().to_vec(); + *cols[idx].mut_metadata() = target_metadata; + + let mut builder = SchemaBuilder::try_from_columns(cols) + .with_context(|_| error::SchemaBuildSnafu { + msg: format!("Failed to convert column schemas into schema for table {table_name}"), + })? + .version(table_schema.version() + 1); + + for (k, v) in table_schema.metadata().iter() { + builder = builder.add_metadata(k, v); + } + + let new_schema = builder.build().with_context(|_| error::SchemaBuildSnafu { + msg: format!( + "Table {table_name} cannot change JSON settings for column {}", + req.column_name + ), + })?; + + let mut meta_builder = self.new_meta_builder(); + let _ = meta_builder + .schema(Arc::new(new_schema)) + .primary_key_indices(self.primary_key_indices.clone()); + + Ok(meta_builder) + } + /// Split requests into different groups using column location info. fn split_requests_by_column_location<'a>( &self, @@ -1603,9 +1676,12 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use datatypes::data_type::ConcreteDataType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; use datatypes::schema::{ ColumnSchema, FulltextAnalyzer, FulltextBackend, Schema, SchemaBuilder, }; + use datatypes::types::JsonType; use super::*; use crate::Error; @@ -1700,6 +1776,90 @@ mod tests { builder.build().unwrap() } + #[test] + fn test_set_json_settings() { + let schema = Arc::new( + SchemaBuilder::try_from_columns(vec![ + ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + json2_column_schema_v1("payload", JsonSettings::new_v2()), + ]) + .unwrap() + .version(123) + .build() + .unwrap(), + ); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + + let alter_kind = AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: "payload".to_string(), + settings: JsonSettings::try_new(vec![], Some(10)).unwrap(), + }, + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + + let payload = new_meta.schema.column_schema_by_name("payload").unwrap(); + let json_metadata: JsonMetadata = + serde_json::from_str(payload.metadata().get("ARROW:extension:metadata").unwrap()) + .unwrap(); + assert!(json_metadata.is_version_2()); + assert_eq!( + Some(10), + json_metadata.json_settings().max_auto_expanded_paths() + ); + assert_eq!(124, new_meta.schema.version()); + assert_eq!(meta.primary_key_indices, new_meta.primary_key_indices); + } + + fn json2_column_schema_v1(name: &str, settings: JsonSettings) -> ColumnSchema { + let mut column_schema = + ColumnSchema::new(name, ConcreteDataType::Json(JsonType::null()), true); + column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new( + JsonMetadata::new_v1(settings), + ))); + column_schema + } + + #[test] + fn test_set_json_settings_rejects_non_json2_column() { + let schema = Arc::new(new_test_schema()); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + let alter_kind = AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: "col2".to_string(), + settings: JsonSettings::new_v2(), + }, + }; + + let err = match meta.builder_with_alter_kind("my_table", &alter_kind) { + Ok(_) => panic!("expected modifying JSON settings on a non-JSON2 column to fail"), + Err(err) => err, + }; + assert!(err.to_string().contains("is not a JSON2 column")); + } + #[test] fn test_modify_time_index_column_type() { let schema = Arc::new(new_test_schema()); diff --git a/src/table/src/predicate.rs b/src/table/src/predicate.rs index 7290c35c70b..8dc85b610f4 100644 --- a/src/table/src/predicate.rs +++ b/src/table/src/predicate.rs @@ -505,13 +505,10 @@ fn extract_from_in_list_expr( for expr in list { if let Expr::Literal(scalar, _) = expr { return_none_if_utf8!(scalar); - if let Some(timestamp) = scalar_value_to_timestamp(scalar, None) { - init_range = init_range.or(&TimestampRange::single(timestamp)) - } else { - // TODO(hl): maybe we should raise an error here since cannot parse - // timestamp value from in list expr - return None; - } + // TODO(hl): maybe we should raise an error here since cannot parse + // timestamp value from in list expr + let timestamp = scalar_value_to_timestamp(scalar, None)?; + init_range = init_range.or(&TimestampRange::single(timestamp)); } } Some(init_range) diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 1ca99268d67..81f396ec703 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -25,11 +25,12 @@ use common_query::AddColumnLocation; use common_time::TimeToLive; use common_time::range::TimestampRange; use datatypes::data_type::ConcreteDataType; +use datatypes::json::JsonSettings; use datatypes::prelude::VectorRef; use datatypes::schema::{ ColumnDefaultConstraint, ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions, }; -use greptime_proto::v1::region::compact_request; +use greptime_proto::v1::region::{build_index_request, compact_request}; use once_cell::sync::Lazy; use serde::{Deserialize, Serialize}; use store_api::metric_engine_consts::{ @@ -48,7 +49,6 @@ use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use crate::error::{ConflictingTableOptionsSnafu, ParseTableOptionSnafu, Result}; use crate::metadata::{TableId, TableVersion}; -use crate::table_reference::TableReference; mod semantic; pub use semantic::*; @@ -371,6 +371,13 @@ pub struct ModifyColumnTypeRequest { pub target_type: ConcreteDataType, } +/// Set JSON2 settings request. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SetJsonSettingsRequest { + pub column_name: String, + pub settings: JsonSettings, +} + /// A family of annotation table options: pure metadata markers that no region /// consumes. Setting or unsetting them only rewrites the table's /// `extra_options`, so the alter skips region dispatch entirely. @@ -661,6 +668,9 @@ pub enum AlterKind { ModifyColumnTypes { columns: Vec, }, + SetJsonSettings { + request: SetJsonSettingsRequest, + }, RenameTable { new_table_name: String, }, @@ -794,6 +804,8 @@ pub struct FlushTableRequest { #[derive(Debug, Clone, Default)] pub struct BuildIndexTableRequest { + /// The index build mode. Absent options select SST indexes. + pub options: Option, pub catalog_name: String, pub schema_name: String, pub table_name: String, @@ -822,25 +834,6 @@ impl Default for CompactTableRequest { } } -/// Truncate table request -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TruncateTableRequest { - pub catalog_name: String, - pub schema_name: String, - pub table_name: String, - pub table_id: TableId, -} - -impl TruncateTableRequest { - pub fn table_ref(&self) -> TableReference<'_> { - TableReference { - catalog: &self.catalog_name, - schema: &self.schema_name, - table: &self.table_name, - } - } -} - #[derive(Debug, Clone, Default, Deserialize, Serialize)] pub struct CopyDatabaseRequest { pub catalog_name: String, diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 16c187d38b6..5bbd8c3c6d0 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -29,7 +29,7 @@ use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datatypes::types::json_type::JsonNativeType; use snafu::ResultExt; -use store_api::storage::{ScanRequest, VectorSearchRequest}; +use store_api::storage::ScanRequest; use crate::error::TablesRecordBatchSnafu; use crate::table::{TableRef, TableType}; @@ -61,18 +61,10 @@ impl DfTableProviderAdapter { self.scan_req.lock().unwrap().output_ordering = Some(order_opts.to_vec()); } - pub fn with_vector_search_hint(&self, hint: VectorSearchRequest) { - self.scan_req.lock().unwrap().vector_search = Some(hint); - } - pub fn with_json_type_hint(&self, hint: std::collections::HashMap) { self.scan_req.lock().unwrap().json_type_hint = hint; } - pub fn get_vector_search_hint(&self) -> Option { - self.scan_req.lock().unwrap().vector_search.clone() - } - #[cfg(feature = "testing")] pub fn get_scan_req(&self) -> ScanRequest { self.scan_req.lock().unwrap().clone() diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 934f838a33a..2db051a1e9c 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -144,8 +144,8 @@ impl RegionScanExec { )) }) .collect::>(); - let ts_col: Option = try { - PhysicalSortExpr::new( + let ts_col: Option = (|| { + Some(PhysicalSortExpr::new( Arc::new( Column::new_with_schema( &metadata.time_index_column().column_schema.name, @@ -157,8 +157,8 @@ impl RegionScanExec { descending: false, nulls_first: true, }, - ) - }; + )) + })(); let eq_props = match request.distribution { Some(TimeSeriesDistribution::PerSeries) => { diff --git a/src/table/src/test_util/memtable.rs b/src/table/src/test_util/memtable.rs index 0ae751d35f4..a0824facd0d 100644 --- a/src/table/src/test_util/memtable.rs +++ b/src/table/src/test_util/memtable.rs @@ -188,7 +188,6 @@ impl Stream for MemtableStream { #[cfg(test)] mod test { use common_recordbatch::util; - use datatypes::prelude::*; use datatypes::schema::ColumnSchema; use datatypes::vectors::{Helper, Int32Vector, StringVector}; diff --git a/tests-fuzz/Cargo.toml b/tests-fuzz/Cargo.toml index 2cd66b8afe5..1189f53c42f 100644 --- a/tests-fuzz/Cargo.toml +++ b/tests-fuzz/Cargo.toml @@ -188,3 +188,12 @@ path = "utils/kafka_wal_helper.rs" test = false bench = false doc = false + +[package.metadata.cargo-shear] +# `common-error` is required by `common-macro`'s `stack_trace_debug`: the +# generated code references `::common_error` paths, so it must stay a direct +# dependency. +# `k8s-openapi` is a feature-pinning dependency: it enables the `v1_30` +# version feature required by the transitive `kube` -> `k8s-openapi` chain, +# but is never referenced by name in source code. +ignored = ["common-error", "k8s-openapi"] diff --git a/tests-fuzz/README.md b/tests-fuzz/README.md index cc9d7eb84ee..7945bae1583 100644 --- a/tests-fuzz/README.md +++ b/tests-fuzz/README.md @@ -6,6 +6,12 @@ cargo install cargo-fuzz ``` +Note: `cargo-fuzz` instruments targets with `-Zsanitizer=fuzzer`, so fuzz +targets must be built with a **nightly** toolchain (the workspace default +toolchain does not need to be nightly — `make fuzz` / `make fuzz-ls` +invoke `cargo +nightly` for you; CI pins its own nightly in +`FUZZ_RUST_TOOLCHAIN`). + 2. Start GreptimeDB 3. Copy the `.env.example`, which is at project root, to `.env` and change the values on need. diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 7a7d0cba72a..79df85475d9 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -11,7 +11,6 @@ path = "tests/main.rs" [features] dashboard = ["servers/dashboard"] -vector_index = [] enterprise = ["frontend/enterprise", "meta-srv/enterprise", "common-meta/enterprise"] [lints] @@ -20,16 +19,11 @@ workspace = true [dependencies] api.workspace = true arrow-flight.workspace = true -async-stream.workspace = true async-trait.workspace = true auth.workspace = true axum.workspace = true -base64.workspace = true cache.workspace = true catalog.workspace = true -chrono.workspace = true -clap.workspace = true -cli.workspace = true client.workspace = true cmd.workspace = true common-base.workspace = true @@ -38,16 +32,12 @@ common-config.workspace = true common-datasource.workspace = true common-error.workspace = true common-event-recorder.workspace = true -common-frontend.workspace = true common-grpc.workspace = true -common-memory-manager.workspace = true common-meta.workspace = true -common-options.workspace = true common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true -common-stat.workspace = true common-telemetry.workspace = true common-test-util.workspace = true common-time.workspace = true @@ -55,63 +45,59 @@ common-wal.workspace = true datanode = { workspace = true } datatypes.workspace = true dotenv.workspace = true -flate2 = "1.0" flow.workspace = true frontend.workspace = true futures.workspace = true futures-util.workspace = true http.workspace = true hyper-util = { workspace = true, features = ["tokio"] } -log-query = { workspace = true } -loki-proto.workspace = true meta-client.workspace = true -meta-srv.workspace = true +meta-srv = { workspace = true, features = ["mock"] } mito2.workspace = true object-store.workspace = true operator.workspace = true -plugins.workspace = true prost.workspace = true query.workspace = true rand.workspace = true serde_json.workspace = true servers.workspace = true session.workspace = true -snafu.workspace = true sql.workspace = true -sqlx = { workspace = true, features = [ - "mysql", - "postgres", - "chrono", -] } standalone.workspace = true store-api.workspace = true substrait.workspace = true table.workspace = true tempfile.workspace = true -time = "0.3" tokio.workspace = true tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true tower.workspace = true uuid.workspace = true -zstd.workspace = true [dev-dependencies] +base64.workspace = true catalog = { workspace = true, features = ["testing"] } +chrono.workspace = true +clap.workspace = true cli.workspace = true client = { workspace = true, features = ["testing"] } +common-frontend.workspace = true +common-memory-manager.workspace = true common-meta = { workspace = true, features = ["testing"] } +common-options.workspace = true common-procedure-test.workspace = true common-wal = { workspace = true, features = ["testing"] } datafusion.workspace = true datafusion-expr.workspace = true either.workspace = true +flate2 = "1.0" frontend = { workspace = true, features = ["testing"] } hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true -meta-srv = { workspace = true, features = ["mock"] } +log-query.workspace = true +loki-proto.workspace = true mito2 = { workspace = true, features = ["test", "test-shared-fs-region-migration"] } mysql_async = { version = "0.37", default-features = false, features = [ "time", @@ -120,9 +106,9 @@ mysql_async = { version = "0.37", default-features = false, features = [ opentelemetry-proto.workspace = true operator = { workspace = true, features = ["testing"] } otel-arrow-rust.workspace = true -partition.workspace = true paste.workspace = true pipeline.workspace = true +plugins.workspace = true prost.workspace = true rand.workspace = true rskafka.workspace = true @@ -130,11 +116,18 @@ rstest.workspace = true rstest_reuse.workspace = true sea-query.workspace = true servers = { workspace = true, features = ["testing"] } -session = { workspace = true, features = ["testing"] } similar-asserts.workspace = true +sqlx = { workspace = true, features = [ + "mysql", + "postgres", + "chrono", +] } store-api.workspace = true +temp-env.workspace = true +time = "0.3" tokio-postgres = { workspace = true } tokio-util.workspace = true url = "2.3" urlencoding = "2.1" yaml-rust = "0.4" +zstd.workspace = true diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 3e53b6104a0..bd357033a7e 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -17,13 +17,15 @@ use std::env; use std::net::TcpListener; use std::ops::RangeInclusive; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; use std::time::Duration; use api::v1::health_check_server::HealthCheckServer; use api::v1::region::region_server::RegionServer; use arrow_flight::flight_service_server::FlightServiceServer; use cache::{ - build_datanode_cache_registry, build_fundamental_cache_registry, + build_datanode_layered_cache_registry, build_fundamental_cache_registry, with_default_composite_cache_registry, }; use catalog::information_extension::DistributedInformationExtension; @@ -36,6 +38,7 @@ use client::client_manager::NodeClients; use cmd::frontend::create_heartbeat_task; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; +use common_event_recorder::EventRecorderOptions; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; @@ -73,9 +76,11 @@ use servers::grpc::{GrpcOptions, HealthCheckHandler}; use servers::server::ServerHandlers; use store_api::storage::RegionId; use tempfile::TempDir; +use tokio::io::{AsyncRead, AsyncWrite}; use tonic::codec::CompressionEncoding; use tonic::transport::Server; -use tower::service_fn; +use tonic::transport::server::Connected; +use tower::{Layer, Service, service_fn}; use uuid::Uuid; use crate::test_util::{ @@ -83,11 +88,31 @@ use crate::test_util::{ create_tmp_dir_and_datanode_opts, }; +/// Flight `DoGet` requests another node sent to a datanode over its network address (see +/// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). +/// +/// The frontend and datanodes inside the test process use in-process clients, so counted requests +/// come from remote nodes querying regions owned by this datanode. +#[derive(Debug, Default)] +pub struct DatanodeRpcStats { + do_get_requests: AtomicUsize, +} + +impl DatanodeRpcStats { + /// The number of Flight `DoGet` requests another node sent to the datanode. + pub fn requests(&self) -> usize { + self.do_get_requests.load(Ordering::Relaxed) + } +} + pub struct GreptimeDbCluster { pub guards: Vec, pub datanode_options: Vec, pub datanode_instances: HashMap, + /// The stats of the network address of every datanode, empty unless the cluster serves the + /// datanodes at their own addresses. + pub datanode_rpc_stats: HashMap>, pub kv_backend: KvBackendRef, pub metasrv: Arc, pub frontend: Arc, @@ -98,6 +123,13 @@ impl GreptimeDbCluster { &self.frontend.instance } + /// The stats of the network address of the datanode `datanode_id`, if the cluster serves the + /// datanodes at their own addresses (see + /// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). + pub fn datanode_rpc_stats(&self, datanode_id: DatanodeId) -> Option<&Arc> { + self.datanode_rpc_stats.get(&datanode_id) + } + /// List all SST files from all datanodes. pub async fn list_sst_files_from_all_datanodes(&self) -> BTreeSet { let mut sst_files = BTreeSet::new(); @@ -173,7 +205,10 @@ pub struct GreptimeDbClusterBuilder { frontend_auto_create_table: bool, shared_home_dir: Option>, meta_selector: Option, + plugins: Plugins, local_file_access: LocalFileAccess, + event_recorder_options: EventRecorderOptions, + real_datanode_grpc_addr: bool, } impl GreptimeDbClusterBuilder { @@ -208,7 +243,10 @@ impl GreptimeDbClusterBuilder { frontend_auto_create_table: true, shared_home_dir: None, meta_selector: None, + plugins: Plugins::default(), local_file_access: LocalFileAccess::default(), + event_recorder_options: EventRecorderOptions::default(), + real_datanode_grpc_addr: false, } } @@ -253,6 +291,15 @@ impl GreptimeDbClusterBuilder { self } + #[must_use] + pub fn with_event_recorder_options( + mut self, + event_recorder_options: EventRecorderOptions, + ) -> Self { + self.event_recorder_options = event_recorder_options; + self + } + /// Sets whether the test frontend automatically creates tables on write. #[must_use] pub fn with_frontend_auto_create_table(mut self, auto_create_table: bool) -> Self { @@ -272,12 +319,34 @@ impl GreptimeDbClusterBuilder { self } + /// Sets the [Plugins] used by the metasrv. + #[must_use] + pub fn with_plugins(mut self, plugins: Plugins) -> Self { + self.plugins = plugins; + self + } + /// Configure the frontend COPY sandbox for filesystem integration tests. pub fn with_local_file_access(mut self, access: LocalFileAccess) -> Self { self.local_file_access = access; self } + /// Serves every datanode on its own real localhost TCP port and registers that + /// address in metasrv, instead of the default in-process duplex placeholder address + /// that all datanodes share. + /// + /// Datanodes serve the region queries of the plans they execute by querying the + /// leader of every region of the plan (see `DatanodeRegionQueryHandler`), so a test + /// that exercises a datanode-to-datanode region query needs each datanode to be + /// reachable at the address recorded in the table route, i.e. the address the + /// datanode registers with its heartbeat. + #[must_use] + pub fn with_real_datanode_grpc_addr(mut self, real_datanode_grpc_addr: bool) -> Self { + self.real_datanode_grpc_addr = real_datanode_grpc_addr; + self + } + pub async fn build_with( &self, datanode_options: Vec, @@ -285,6 +354,7 @@ impl GreptimeDbClusterBuilder { guards: Vec, ) -> GreptimeDbCluster { let datanodes = datanode_options.len(); + let channel_config = ChannelConfig::new().timeout(Some(Duration::from_secs(20))); let datanode_clients = Arc::new(NodeClients::new(channel_config)); @@ -303,18 +373,20 @@ impl GreptimeDbClusterBuilder { ..Default::default() }, gc: self.metasrv_gc_config.clone(), + event_recorder: self.event_recorder_options.clone(), ..Default::default() }; test_util::prepare_another_catalog_and_schema_with_kv_backend(self.kv_backend.clone()) .await; - let metasrv = meta_srv::mocks::mock( + let metasrv = meta_srv::mocks::mock_with_plugins( opt, self.kv_backend.clone(), self.meta_selector.clone(), Some(datanode_clients.clone()), None, + self.plugins.clone(), ) .await; @@ -322,7 +394,22 @@ impl GreptimeDbClusterBuilder { .build_datanodes_with_options(&metasrv, &datanode_options) .await; - build_datanode_clients(datanode_clients.clone(), &datanode_instances).await; + let datanode_rpc_stats = if self.real_datanode_grpc_addr { + datanode_options + .iter() + .map(|opts| (opts.node_id.unwrap(), Arc::new(DatanodeRpcStats::default()))) + .collect::>() + } else { + HashMap::new() + }; + + build_datanode_clients( + datanode_clients.clone(), + &datanode_instances, + &datanode_options, + &datanode_rpc_stats, + ) + .await; self.wait_datanodes_alive(metasrv.metasrv.meta_peer_client(), datanodes) .await; @@ -341,6 +428,7 @@ impl GreptimeDbClusterBuilder { datanode_options, guards, datanode_instances, + datanode_rpc_stats, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -367,7 +455,7 @@ impl GreptimeDbClusterBuilder { let home_dir = if let Some(home_dir) = &self.shared_home_dir { home_dir.path().to_str().unwrap().to_string() } else { - let home_tmp_dir = create_temp_dir(&format!("gt_home_{}", &self.cluster_name)); + let home_tmp_dir = create_temp_dir(&format!("gt_home_{}", self.cluster_name)); let home_dir = home_tmp_dir.path().to_str().unwrap().to_string(); guards.push(TestGuard { home_guard: FileDirGuard::new(home_tmp_dir), @@ -397,12 +485,28 @@ impl GreptimeDbClusterBuilder { opts }; opts.node_id = Some(datanode_id); - + if self.real_datanode_grpc_addr { + // Give every datanode its own address, so that the other datanodes can + // reach it (see `with_real_datanode_grpc_addr`). + let addr = self.choose_datanode_grpc_addr(); + opts.grpc.bind_addr = addr.clone(); + opts.grpc.server_addr = addr; + } options.push(opts); } (options, guards) } + /// Chooses a free localhost port for a datanode. + /// + /// Uses a port range disjoint from the one of the frontend servers + /// (see [`Self::build_frontend_options`]). + fn choose_datanode_grpc_addr(&self) -> String { + let localhost = "127.0.0.1"; + let port = self.choose_random_unused_port(24000..=34000, 20, localhost); + format!("{}:{}", localhost, port) + } + async fn build_datanodes_with_options( &self, metasrv: &MockInfo, @@ -412,7 +516,8 @@ impl GreptimeDbClusterBuilder { for opts in options { let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await; - instances.insert(opts.node_id.unwrap(), datanode); + let datanode_id = opts.node_id.unwrap(); + instances.insert(datanode_id, datanode); } instances @@ -450,15 +555,12 @@ impl GreptimeDbClusterBuilder { let meta_backend = new_read_only_meta_kv_backend(meta_client.clone()); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(meta_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(meta_backend.clone())); let mut builder = DatanodeBuilder::new(opts.clone(), Plugins::default(), meta_backend); builder - .with_cache_registry(layered_cache_registry) + .with_cache_registry(layered_cache_registry.clone()) .with_meta_client(meta_client); let mut datanode = builder.build().await.unwrap(); @@ -550,6 +652,7 @@ impl GreptimeDbClusterBuilder { fn build_frontend_options(&self) -> FrontendOptions { let mut fe_opts = FrontendOptions { auto_create_table: self.frontend_auto_create_table, + event_recorder: self.event_recorder_options.clone(), ..Default::default() }; @@ -605,18 +708,34 @@ impl GreptimeDbClusterBuilder { async fn build_datanode_clients( clients: Arc, instances: &HashMap, + options: &[DatanodeOptions], + rpc_stats: &HashMap>, ) { - for (&datanode_id, instance) in instances { - let (addr, client) = create_datanode_client(instance).await; + for opts in options { + let datanode_id = opts.node_id.unwrap(); + let instance = &instances[&datanode_id]; + let (addr, client) = create_datanode_client( + instance, + &opts.grpc.server_addr, + rpc_stats.get(&datanode_id).cloned(), + ) + .await; clients .insert_client(Peer::new(datanode_id, addr), client) .await; } } -async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { - let (client, server) = tokio::io::duplex(1024); - +/// Runs a gRPC server that serves the region server of `datanode` over `incoming` until the +/// cluster is dropped. +fn serve_datanode_region_server( + datanode: &Datanode, + incoming: I, + rpc_stats: Option>, +) where + I: futures::Stream> + Send + 'static, + IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static, +{ let runtime = RuntimeBuilder::default() .worker_threads(2) .thread_name("grpc-handlers") @@ -627,8 +746,9 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { let region_server_handler = RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime); - let _handle = tokio::spawn(async move { + tokio::spawn(async move { Server::builder() + .layer(CountRequestsLayer { rpc_stats }) .add_service( FlightServiceServer::new(flight_handler) .accept_compressed(CompressionEncoding::Gzip) @@ -644,12 +764,40 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { .send_compressed(CompressionEncoding::Zstd), ) .add_service(HealthCheckServer::new(HealthCheckHandler)) - .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)])) + .serve_with_incoming(incoming) .await }); +} + +/// Serves the region server of `datanode` and returns the address and the client of the address. +/// +/// The returned client is always an in-process client: the frontend and the datanodes of the test +/// process reach the region server through a duplex stream, no matter which address they use. This +/// keeps the in-process traffic out of `rpc_stats`. +/// +/// [`test_util::PEER_PLACEHOLDER_ADDR`] is not a real address, so it cannot be served over the +/// network at all. Any other address is a real one: the region server is additionally served on +/// its local TCP port (`rpc_stats` counts the requests of the remote nodes that reach it, e.g. the +/// region query that a datanode sends for the regions of another datanode). +async fn create_datanode_client( + datanode: &Datanode, + addr: &str, + rpc_stats: Option>, +) -> (String, Client) { + let (client, server) = tokio::io::duplex(1024); + serve_datanode_region_server( + datanode, + futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]), + None, + ); + + if addr != test_util::PEER_PLACEHOLDER_ADDR { + let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); + let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); + serve_datanode_region_server(datanode, incoming, rpc_stats); + } let mut client = Some(client); - let addr = test_util::PEER_PLACEHOLDER_ADDR; let channel_manager = ChannelManager::new(); channel_manager .reset_with_connector( @@ -668,12 +816,64 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { ( addr.to_string(), - Client::with_manager_and_urls(channel_manager, [addr]), + // The mock connector pool lives on this single manager, so both lanes + // intentionally share it (same semantics as the legacy constructor). + Client::with_query_and_control_managers(channel_manager.clone(), channel_manager, [addr]), ) } +/// Counts Flight `DoGet` requests received by the wrapped server. +#[derive(Clone)] +struct CountRequests { + inner: S, + rpc_stats: Option>, +} + +impl Service> for CountRequests +where + S: Service>, +{ + type Response = S::Response; + type Error = S::Error; + type Future = S::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + if request.uri().path() == "/arrow.flight.protocol.FlightService/DoGet" + && let Some(rpc_stats) = &self.rpc_stats + { + rpc_stats.do_get_requests.fetch_add(1, Ordering::Relaxed); + } + + self.inner.call(request) + } +} + +#[derive(Clone)] +struct CountRequestsLayer { + rpc_stats: Option>, +} + +impl Layer for CountRequestsLayer { + type Service = CountRequests; + + fn layer(&self, inner: S) -> Self::Service { + CountRequests { + inner, + rpc_stats: self.rpc_stats.clone(), + } + } +} + +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod tests { + use std::sync::atomic::{AtomicBool, Ordering}; + use api::v1::flow::FlowRequest; use api::v1::region::{ ListMetadataRequest, RegionRequest, RegionRequestHeader, region_request, @@ -682,12 +882,14 @@ mod tests { use client::Client; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; + use common_meta::ddl_manager::{DdlManager, DdlManagerConfigurator, DdlManagerConfiguratorRef}; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::node_manager::{DatanodeManager, FlownodeManager}; use common_meta::peer::Peer; use flow::{FlownodeBuilder, FlownodeOptions, FlownodeServiceBuilder, FrontendClient}; + use meta_srv::metasrv::builder::DdlManagerConfigureContext; use super::*; @@ -864,6 +1066,38 @@ mod tests { flownode } + struct TestDdlManagerConfigurator(Arc); + + #[async_trait::async_trait] + impl DdlManagerConfigurator for TestDdlManagerConfigurator { + async fn configure( + &self, + ddl_manager: DdlManager, + _ctx: DdlManagerConfigureContext, + ) -> std::result::Result { + self.0.store(true, Ordering::SeqCst); + Ok(ddl_manager) + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_build_cluster_with_plugins() { + let configured = Arc::new(AtomicBool::new(false)); + let plugins = Plugins::default(); + plugins.insert::>(Arc::new( + TestDdlManagerConfigurator(configured.clone()), + )); + + let _cluster = GreptimeDbClusterBuilder::new("test_cluster_with_plugins") + .await + .with_datanodes(1) + .with_plugins(plugins) + .build(false) + .await; + + assert!(configured.load(Ordering::SeqCst)); + } + #[tokio::test(flavor = "multi_thread")] async fn test_region_requester_health_check_and_list_metadata() { // Arrange: configure active and inactive synthetic peers over the duplex harness. @@ -928,13 +1162,41 @@ mod tests { // Assert: the inactive peer was rejected and the production Flow service handled the RPC. assert_eq!( - (vec![active_peer.addr], vec![inactive_peer.addr]), + (vec![active_peer.addr.clone()], vec![inactive_peer.addr]), client.peer_addresses_by_state() ); assert_eq!( StatusCode::InvalidArguments, response.unwrap_err().status_code() ); + + for value in ["", "unknown", "follower"] { + let mut query_context = + api::v1::QueryContext::from(session::context::QueryContext::arc().as_ref()); + query_context.extensions.insert( + session::hints::READ_PREFERENCE_EXTENSION_KEY.to_string(), + value.to_string(), + ); + let error = clients + .flownode(&active_peer) + .await + .handle(FlowRequest { + header: Some(api::v1::flow::FlowRequestHeader { + tracing_context: Default::default(), + query_context: Some(query_context), + }), + body: None, + }) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + let message = error.output_msg(); + assert!(message.contains("read_preference"), "{message}"); + if !value.is_empty() { + assert!(message.contains(value), "{message}"); + } + } + flownode.shutdown().await.unwrap(); } } diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index fb27a273516..e74d4d5bf70 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod test { use std::collections::HashMap; use std::net::SocketAddr; @@ -45,6 +47,9 @@ mod test { RecordBatch, RecordBatchStreamWrapper, RecordBatches, SendableRecordBatchStream, }; use common_telemetry::tracing_context::TracingContext; + use datatypes::arrow::datatypes::{ + DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit, + }; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; @@ -64,14 +69,17 @@ mod test { use session::context::QueryContextRef; use session::hints::INSERT_SKIP_WAL_HINT; use tokio::net::TcpListener; - use tokio_stream::wrappers::TcpListenerStream; + use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream}; use tonic::transport::Server as TonicServer; use tonic::{Response, Status}; use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; use crate::grpc::query_and_expect; - use crate::test_util::{MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server}; + use crate::test_util::{ + MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server, + setup_grpc_server_with_auto_create_table_disabled, + }; use crate::tests::test_util::MockInstance; struct SlowFlightCraft; @@ -465,6 +473,315 @@ mod test { db.shutdown().await; } + #[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_auto_add_columns( + #[case] distributed: bool, + #[case] partitioned: bool, + ) { + let mut db = MockInstanceImpl::new( + &format!("flight_bulk_auto_add_columns_{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, PRIMARY KEY (a)){partition}" + )).await.unwrap(); + // Existing rows require the inferred column to be nullable. + client + .sql("INSERT INTO foo VALUES (100, 100)") + .await + .unwrap(); + test_put_record_batches(&client, create_record_batches(-4)).await; + + query_and_expect( + db.frontend().as_ref(), + "SELECT a, \"B\" FROM foo ORDER BY a", + "\ ++-----+-----+ +| a | B | ++-----+-----+ +| -4 | s4 | +| -3 | s3 | +| -2 | s2 | +| -1 | s1 | +| 0 | s0 | +| 1 | s-1 | +| 2 | s-2 | +| 3 | s-3 | +| 4 | s-4 | +| 100 | | ++-----+-----+", + ) + .await; + server.shutdown().await.unwrap(); + db.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns_once_per_stream() { + let (db, server) = setup_grpc_server( + StorageType::File, + "test_flight_bulk_auto_add_columns_once_per_stream", + ) + .await; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a))") + .await + .unwrap(); + let batches = create_record_batches(1); + let schema = batches[0].schema.arrow_schema().clone(); + let (tx, rx) = tokio::sync::mpsc::channel(4); + // The client must be able to receive the handshake before sending its schema. + let mut responses = client + .do_put(ReceiverStream::new(rx).boxed()) + .await + .unwrap(); + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + tx.send(encode_put_schema(schema.as_ref())).await.unwrap(); + + let empty_batch = common_recordbatch::DfRecordBatch::new_empty(schema); + for round in 0..2 { + for message in + FlightEncoder::default().encode(FlightMessage::RecordBatch(empty_batch.clone())) + { + tx.send(message).await.unwrap(); + } + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + if round == 0 { + // Even an empty first batch initializes the schema. Dropping the + // column makes a repeated reconciliation on the next batch observable. + client + .sql("ALTER TABLE foo DROP COLUMN \"B\"") + .await + .unwrap(); + } + } + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| a | +| ts | ++-------------+", + ) + .await; + drop(tx); + assert!(responses.next().await.is_none()); + + // A new stream reconciles again and persists the new column's values. + test_put_record_batches(&client, batches).await; + query_and_expect( + db.frontend().as_ref(), + "SELECT count(\"B\") AS n FROM foo", + "\ ++---+ +| n | ++---+ +| 9 | ++---+", + ) + .await; + server.shutdown().await.unwrap(); + } + + #[rstest] + #[case::hint(false)] + #[case::config(true)] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns_disabled(#[case] disabled_by_config: bool) { + let name = "test_flight_bulk_auto_add_columns_disabled"; + let (db, server) = if disabled_by_config { + setup_grpc_server_with_auto_create_table_disabled(StorageType::File, name).await + } else { + setup_grpc_server(StorageType::File, name).await + }; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a))") + .await + .unwrap(); + let batch = create_record_batches(1).remove(0).into_df_record_batch(); + let schema = batch.schema(); + let mut messages = vec![encode_put_schema(schema.as_ref())]; + messages.extend(FlightEncoder::default().encode(FlightMessage::RecordBatch( + common_recordbatch::DfRecordBatch::new_empty(schema), + ))); + messages.extend(FlightEncoder::default().encode(FlightMessage::RecordBatch(batch))); + // A request hint cannot override the server-side setting. + let hint = if disabled_by_config { "true" } else { "false" }; + let mut responses = client + .do_put_with_hints( + tokio_stream::iter(messages).boxed(), + &[("auto_create_table", hint)], + ) + .await + .unwrap(); + // The handshake and empty batch succeed, but the non-empty write must fail. + for _ in 0..2 { + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + } + let Some(Err(err)) = responses.next().await else { + panic!("expected the bulk write with an unknown column to fail"); + }; + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + let message = err.to_string(); + assert!(message.contains("Column 'B' not found"), "{message}"); + assert!(responses.next().await.is_none()); + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| a | +| ts | ++-------------+", + ) + .await; + query_and_expect( + db.frontend().as_ref(), + "SELECT count(*) AS n FROM foo", + "\ ++---+ +| n | ++---+ +| 0 | ++---+", + ) + .await; + test_put_record_batches_with_hints( + &client, + create_record_batches_without_nullable_column(1), + &[("auto_create_table", hint)], + ) + .await; + server.shutdown().await.unwrap(); + } + + #[rstest] + #[case::list(ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Int32, true))))] + #[case::list_unsupported_child(ArrowDataType::List(Arc::new(Field::new( + "item", + ArrowDataType::FixedSizeBinary(16), + true, + ))))] + #[case::struct_unsupported_child(ArrowDataType::Struct( + vec![Field::new("item", ArrowDataType::FixedSizeBinary(16), true)].into(), + ))] + #[case::dictionary(ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ))] + #[case::dictionary_unsupported_value(ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::FixedSizeBinary(16)), + ))] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_rejects_nested_auto_add_without_altering_table( + #[case] data_type: ArrowDataType, + ) { + let (db, server) = setup_grpc_server( + StorageType::File, + "test_flight_bulk_rejects_nested_auto_add_without_altering_table", + ) + .await; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX)") + .await + .unwrap(); + + // Use raw Arrow fields: Greptime's schema conversion itself used to panic + // on unsupported child types. A preceding scalar must not be partially added. + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("new_scalar", ArrowDataType::Int32, true), + Field::new("new_nested", data_type.clone(), true), + ])); + let mut encoder = FlightEncoder::default(); + let mut schema_data = encoder.encode_schema(schema.as_ref()); + schema_data.flight_descriptor = Some(FlightDescriptor { + r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, + path: vec!["foo".to_string()], + ..Default::default() + }); + let mut messages = vec![schema_data]; + // An empty first batch still performs schema reconciliation. + messages.extend(encoder.encode(FlightMessage::RecordBatch( + common_recordbatch::DfRecordBatch::new_empty(schema), + ))); + let mut responses = client + .do_put(tokio_stream::iter(messages).boxed()) + .await + .unwrap(); + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + let Some(Err(err)) = responses.next().await else { + panic!("expected an unsupported bulk schema error"); + }; + assert_eq!(err.status_code(), StatusCode::Unsupported); + let message = err.to_string(); + assert!(message.contains("new_nested"), "{message}"); + assert!(message.contains(&format!("{data_type:?}")), "{message}"); + assert!(responses.next().await.is_none()); + + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| ts | ++-------------+", + ) + .await; + server.shutdown().await.unwrap(); + } + #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -673,7 +990,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); - assert!(result.metrics.is_ready()); + result.metrics.wait_ready().await; assert!(result.region_watermark_map().is_none()); let err = client @@ -699,6 +1016,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); + result.metrics.wait_ready().await; assert_eq!( result.region_watermark_map(), Some(std::collections::HashMap::from([previous_watermark])) @@ -814,34 +1132,23 @@ mod test { let requests_count = record_batches.len(); let schema = record_batches[0].schema.arrow_schema().clone(); - let stream = futures::stream::once(async move { - let mut schema_data = FlightEncoder::default().encode_schema(schema.as_ref()); - let metadata = DoPutMetadata::new(0); - schema_data.app_metadata = serde_json::to_vec(&metadata).unwrap().into(); - // first message in "DoPut" stream should carry table name in flight descriptor - schema_data.flight_descriptor = Some(FlightDescriptor { - r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, - path: vec!["foo".to_string()], - ..Default::default() - }); - schema_data - }) - .chain( - tokio_stream::iter(record_batches) - .enumerate() - .flat_map(|(i, x)| { - let mut encoder = FlightEncoder::default(); - let message = FlightMessage::RecordBatch(x.into_df_record_batch()); - let mut data = encoder.encode(message); - let metadata = DoPutMetadata::new((i + 1) as i64); - data.iter_mut().for_each(|x| { - x.app_metadata = serde_json::to_vec(&metadata).unwrap().into() - }); - tokio_stream::iter(data) - }) - .boxed(), - ) - .boxed(); + let stream = futures::stream::once(async move { encode_put_schema(schema.as_ref()) }) + .chain( + tokio_stream::iter(record_batches) + .enumerate() + .flat_map(|(i, x)| { + let mut encoder = FlightEncoder::default(); + let message = FlightMessage::RecordBatch(x.into_df_record_batch()); + let mut data = encoder.encode(message); + let metadata = DoPutMetadata::new((i + 1) as i64); + data.iter_mut().for_each(|x| { + x.app_metadata = serde_json::to_vec(&metadata).unwrap().into() + }); + tokio_stream::iter(data) + }) + .boxed(), + ) + .boxed(); let response_stream = client.do_put_with_hints(stream, hints).await.unwrap(); @@ -861,6 +1168,17 @@ mod test { assert_eq!(requests_count + 1, responses_count); } + fn encode_put_schema(schema: &datatypes::arrow::datatypes::Schema) -> FlightData { + let mut schema_data = FlightEncoder::default().encode_schema(schema); + schema_data.app_metadata = serde_json::to_vec(&DoPutMetadata::new(0)).unwrap().into(); + schema_data.flight_descriptor = Some(FlightDescriptor { + r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, + path: vec!["foo".to_string()], + ..Default::default() + }); + schema_data + } + fn create_record_batches_without_nullable_column(start: i64) -> Vec { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( diff --git a/tests-integration/src/grpc/network.rs b/tests-integration/src/grpc/network.rs index 427317fb05d..97e4c56494a 100644 --- a/tests-integration/src/grpc/network.rs +++ b/tests-integration/src/grpc/network.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod tests { use std::io; use std::pin::Pin; diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index dacab7ad6bd..3e447bb0194 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -14,12 +14,28 @@ #[cfg(test)] mod test { - use std::sync::Arc; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + use TimeUnit as ArrowTimeUnit; + use api::region::RegionResponse; + use api::v1::region::{RegionRequest, region_request}; use client::{DEFAULT_CATALOG_NAME, OutputData}; - use common_recordbatch::RecordBatches; + use common_meta::cache::LayeredCacheRegistryBuilder; + use common_meta::node_manager::{ + Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef, NodeManagerRef, + }; + use common_meta::peer::Peer; + use common_query::request::QueryRequest; + use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::{ + DataType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, + }; use frontend::instance::Instance; + use frontend::instance::builder::FrontendBuilder; + use frontend::service_config::BatcherOptions; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::any_value::Value as Val; @@ -34,13 +50,574 @@ mod test { use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; use pipeline::{GreptimePipelineParams, PipelineWay}; use serde_json::json; + use servers::batcher::BatchingProtocol; + use servers::otlp::trace::{TraceAuxData, span}; use servers::query_handler::OpenTelemetryProtocolHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + use store_api::region_engine::RegionEngine; + use store_api::storage::RegionId; + use tokio::sync::oneshot; - use crate::standalone::GreptimeDbStandaloneBuilder; + use crate::standalone::{GreptimeDbStandalone, GreptimeDbStandaloneBuilder}; use crate::tests; + #[tokio::test(flavor = "multi_thread")] + async fn test_trace_aux_cache() { + check_trace_aux_cache(false).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_trace_aux_cache_with_batching_v0_v1() { + check_trace_aux_cache(true).await; + } + + #[test] + fn test_trace_aux_async_completion() { + temp_env::with_var("PENDING_ROWS_BATCH_SYNC", Some("false"), || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(check_trace_aux_async_completion()); + }); + } + + #[test] + fn test_trace_aux_admission() { + temp_env::with_var("PENDING_ROWS_BATCH_SYNC", Some("false"), || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + tokio::time::timeout(Duration::from_secs(60), check_trace_aux_admission()) + .await + .expect("auxiliary admission must make progress"); + }); + }); + } + + struct TraceAuxWriteGate { + region_id: RegionId, + started: oneshot::Sender<()>, + release: oneshot::Receiver, + } + + struct GatedTraceNodeManager { + inner: NodeManagerRef, + gate: Arc>>, + } + + #[async_trait::async_trait] + impl DatanodeManager for GatedTraceNodeManager { + async fn datanode(&self, peer: &Peer) -> DatanodeRef { + Arc::new(GatedTraceDatanode { + inner: self.inner.datanode(peer).await, + gate: self.gate.clone(), + }) + } + } + + #[async_trait::async_trait] + impl FlownodeManager for GatedTraceNodeManager { + async fn flownode(&self, peer: &Peer) -> FlownodeRef { + self.inner.flownode(peer).await + } + } + + struct GatedTraceDatanode { + inner: DatanodeRef, + gate: Arc>>, + } + + #[async_trait::async_trait] + impl Datanode for GatedTraceDatanode { + async fn handle( + &self, + request: RegionRequest, + ) -> common_meta::error::Result { + let gate = { + let mut gate = self.gate.lock().unwrap(); + if gate.as_ref().is_some_and(|gate| { + matches!(&request.body, Some(region_request::Body::Inserts(inserts)) + if inserts.requests.iter().any(|insert| insert.region_id == gate.region_id.as_u64())) + }) { + gate.take() + } else { + None + } + }; + if let Some(gate) = gate { + gate.started.send(()).unwrap(); + if !gate.release.await.unwrap() { + return common_meta::error::UnexpectedSnafu { + err_msg: "injected auxiliary write failure", + } + .fail(); + } + } + self.inner.handle(request).await + } + + async fn handle_query( + &self, + request: QueryRequest, + ) -> common_meta::error::Result { + self.inner.handle_query(request).await + } + } + + async fn check_trace_aux_admission() { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_admission") + .with_table_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 1, + max_inflight_requests: 1, + ..Default::default() + }) + .build() + .await; + let source = standalone.fe_instance(); + let gate = Arc::new(Mutex::new(None)); + let cache_registry = Arc::new( + cache::with_default_composite_cache_registry( + LayeredCacheRegistryBuilder::default().add_cache_registry( + cache::build_fundamental_cache_registry(standalone.kv_backend.clone()), + ), + ) + .unwrap() + .build(), + ); + let mut options = standalone.opts.frontend_options(); + options.otlp.trace_ingest_chunk_size = 1; + let instance = Arc::new( + FrontendBuilder::new( + options, + standalone.kv_backend.clone(), + cache_registry, + source.catalog_manager().clone(), + Arc::new(GatedTraceNodeManager { + inner: source.node_manager().clone(), + gate: gate.clone(), + }), + source.procedure_executor().clone(), + Arc::new(catalog::process_manager::ProcessManager::new( + "trace_aux_admission".into(), + None, + )), + ) + .try_build() + .await + .unwrap(), + ); + + for version in 0..=1 { + let table_name = format!("trace_aux_admission_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(true); + let ctx = Arc::new(context); + let write = async |request| { + instance + .traces( + instance.clone(), + request, + if version == 0 { + PipelineWay::OtlpTraceDirectV0 + } else { + PipelineWay::OtlpTraceDirectV1 + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx.clone(), + ) + .await + .unwrap() + }; + let cached = |request: ExportTraceServiceRequest| { + let mut data = TraceAuxData::default(); + for span in span::parse(request).iter().flat_map(|group| &group.spans) { + data.observe_span(span); + } + instance.trace_aux_data_cached_for_test(&table_name, data, &ctx) + }; + let warm = trace_aux_request(); + assert_eq!(1, write(warm.clone()).await.accepted_spans); + while !cached(warm.clone()) { + tokio::task::yield_now().await; + } + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}_operations"), + None, + ) + .await + .unwrap() + .unwrap(); + let route = instance + .partition_manager() + .find_physical_table_route(table.table_info().table_id()) + .await + .unwrap(); + let region_id = route.region_routes[0].region.id; + + for success in [true, false] { + let mut cold = warm.clone(); + let spans = &mut cold.resource_spans[0].scope_spans[0].spans; + spans[0].name = format!("cold_{success}"); + spans.push(spans[0].clone()); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + *gate.lock().unwrap() = Some(TraceAuxWriteGate { + region_id, + started: started_tx, + release: release_rx, + }); + + // Two chunks must progress with just one main and one auxiliary slot. + assert_eq!(2, write(cold.clone()).await.accepted_spans); + started_rx.await.unwrap(); + assert!( + !cached(cold.clone()), + "pending writes must not warm the cache" + ); + let before = trace_aux_sequences(&standalone, &table_name).await; + let mut waiting = Box::pin(write(cold.clone())); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + + // Main admission is free while auxiliary storage is blocked. Warm + // requests must bypass auxiliary admission and keep making progress. + assert_eq!(1, write(warm.clone()).await.accepted_spans); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + release_tx.send(success).unwrap(); + assert_eq!(2, waiting.await.accepted_spans); + if !success { + // The failed write cannot satisfy the waiting request's miss. + while !cached(cold.clone()) { + tokio::task::yield_now().await; + } + } + let after = trace_aux_sequences(&standalone, &table_name).await; + assert_eq!(before[1], after[1], "cached service must not be rewritten"); + assert_eq!( + before[2] + 1, + after[2], + "exactly one successful operation write" + ); + } + } + } + + fn trace_aux_request() -> ExportTraceServiceRequest { + serde_json::from_value(json!({ + "resourceSpans": [{ + "resource": {"attributes": [ + {"key": "service.name", "value": {"stringValue": "svc"}} + ]}, + "scopeSpans": [{"spans": [{ + "traceId": "c05d7a4ec8e1f231f02ed6e8da8655b4", + "spanId": "9630f2916e2f7909", + "name": "operation", + "kind": 2, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000" + }]}] + }] + })) + .unwrap() + } + + async fn trace_aux_sequences(standalone: &GreptimeDbStandalone, table_name: &str) -> Vec { + let instance = standalone.fe_instance(); + let mut sequences = Vec::new(); + for suffix in ["", "_services", "_operations"] { + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}{suffix}"), + None, + ) + .await + .unwrap() + .unwrap(); + let route = instance + .partition_manager() + .find_physical_table_route(table.table_info().table_id()) + .await + .unwrap(); + assert_eq!(route.region_routes.len(), 1); + sequences.push( + standalone + .mito_engine + .get_committed_sequence(route.region_routes[0].region.id) + .await + .unwrap(), + ); + } + sequences + } + + async fn check_trace_aux_async_completion() { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_async_completion") + .with_table_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 2, + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + // V2 bulk ingestion has a separate Variant/Struct conversion failure. + for version in 0..=1 { + let table_name = format!("trace_aux_async_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(true); + let ctx = Arc::new(context); + let request = trace_aux_request(); + let write = async |request, ctx| { + let outcome = tokio::time::timeout( + Duration::from_secs(30), + instance.traces( + instance.clone(), + request, + if version == 0 { + PipelineWay::OtlpTraceDirectV0 + } else { + PipelineWay::OtlpTraceDirectV1 + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx, + ), + ) + .await + .expect("async trace ingestion must return before the batch flush") + .unwrap(); + assert_eq!((outcome.accepted_spans, outcome.rejected_spans), (1, 0)); + assert!(outcome.error_message.is_none()); + }; + write(request.clone(), ctx.clone()).await; + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) = 0 FROM {table_name}"), + ctx.clone(), + ) + .await + .unwrap(); + for suffix in ["_services", "_operations"] { + assert!( + instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}{suffix}"), + None + ) + .await + .unwrap() + .is_none() + ); + } + + // The second submission deterministically triggers the main flush. + // Omitting the service avoids a second concurrent auxiliary write. + let mut flush_request = request.clone(); + flush_request.resource_spans[0] + .resource + .as_mut() + .unwrap() + .attributes + .clear(); + write(flush_request, ctx.clone()).await; + tokio::time::timeout(Duration::from_secs(30), async { + loop { + let mut aux_data = TraceAuxData::default(); + for span in span::parse(request.clone()) + .into_iter() + .flat_map(|group| group.spans) + { + aux_data.observe_span(&span); + } + if instance.trace_aux_data_cached_for_test(&table_name, aux_data, &ctx) { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("deferred auxiliary writes must populate the cache"); + for suffix in ["_services", "_operations"] { + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) = 1 FROM {table_name}{suffix}"), + ctx.clone(), + ) + .await + .unwrap(); + } + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) > 0 FROM {table_name}"), + ctx.clone(), + ) + .await + .unwrap(); + + let before = trace_aux_sequences(&standalone, &table_name).await; + // A direct warm request finishes before we compare storage sequences. + let mut warm_ctx = ctx.fork(); + warm_ctx.set_batching_enabled(false); + write(request, Arc::new(warm_ctx)).await; + let after = trace_aux_sequences(&standalone, &table_name).await; + assert!(after[0] > before[0]); + assert_eq!(after[1..], before[1..], "v{version}"); + } + } + + async fn check_trace_aux_cache(batching: bool) { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_cache") + .with_table_batcher(if batching { + BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + // Two main spans trigger a flush, while each auxiliary table + // has only one row and would wait an hour if batched. + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 2, + ..Default::default() + } + } else { + BatcherOptions::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let cloned_instance = Arc::new(instance.as_ref().clone()); + + // V2 bulk writes currently fail on Variant/Struct schema conversion + // before reaching auxiliary ingestion; cover V2 with direct main writes. + let max_version = if batching { 1 } else { 2 }; + for version in 0..=max_version { + let table_name = format!("trace_aux_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(batching); + let ctx = Arc::new(context); + let mut request = trace_aux_request(); + if batching { + let spans = &mut request.resource_spans[0].scope_spans[0].spans; + spans.push(spans[0].clone()); + } + let accepted_spans = if batching { 2 } else { 1 }; + let write = async |frontend: &Arc, request| { + tokio::time::timeout( + Duration::from_secs(30), + frontend.traces( + frontend.clone(), + request, + match version { + 0 => PipelineWay::OtlpTraceDirectV0, + 1 => PipelineWay::OtlpTraceDirectV1, + _ => PipelineWay::OtlpTraceDirectV2, + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx.clone(), + ), + ) + .await + .expect("auxiliary writes must not wait for the batch flush interval") + .unwrap() + }; + + // An incompatible auxiliary column fails after the main span is + // accepted. No candidate key may be cached from the failed batch. + instance + .do_query( + &format!( + "CREATE TABLE {table_name}_operations (\ + \"timestamp\" TIMESTAMP(9) TIME INDEX, service_name STRING, \ + span_name STRING, span_kind BIGINT)" + ), + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + let failed_aux = write(instance, request.clone()).await; + assert_eq!( + (failed_aux.accepted_spans, failed_aux.rejected_spans), + (accepted_spans, 0) + ); + assert!( + failed_aux + .error_message + .unwrap() + .contains("Auxiliary trace tables") + ); + + instance + .do_query(&format!("DROP TABLE {table_name}_operations"), ctx.clone()) + .await + .remove(0) + .unwrap(); + let retry = write(instance, request.clone()).await; + assert_eq!( + (retry.accepted_spans, retry.rejected_spans), + (accepted_spans, 0) + ); + assert!(retry.error_message.is_none(), "{retry:?}"); + let before = trace_aux_sequences(&standalone, &table_name).await; + assert!(before.iter().all(|sequence| *sequence > 0)); + + // Storage sequences, unlike row counts in an upsert table, prove + // the writes were skipped. Cloned Instances must share the cache. + let warm = write(&cloned_instance, request.clone()).await; + assert_eq!( + (warm.accepted_spans, warm.rejected_spans), + (accepted_spans, 0) + ); + assert!(warm.error_message.is_none(), "{warm:?}"); + let after = trace_aux_sequences(&standalone, &table_name).await; + assert!(after[0] > before[0]); + assert_eq!(after[1..], before[1..], "v{version}"); + + request.resource_spans[0].scope_spans[0].spans[0].name = "new_operation".into(); + let new_operation = write(instance, request).await; + assert_eq!( + (new_operation.accepted_spans, new_operation.rejected_spans), + (accepted_spans, 0) + ); + assert!(new_operation.error_message.is_none(), "{new_operation:?}"); + let new_sequences = trace_aux_sequences(&standalone, &table_name).await; + assert!(new_sequences[0] > after[0]); + assert_eq!(new_sequences[1], after[1]); + assert!(new_sequences[2] > after[2]); + } + } + #[tokio::test(flavor = "multi_thread")] async fn test_trace_v2_on_standalone() -> Result<(), Box> { let standalone = GreptimeDbStandaloneBuilder::new("trace_v2_standalone") @@ -503,6 +1080,493 @@ WITH( Ok(()) } + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_logical_batcher_non_millisecond_physical_table() { + // Microsecond and nanosecond physical tables must both use the + // logical batcher: requests are converted to the physical table's + // unit during batch alignment + // (). + run_otlp_logical_batcher_non_millisecond_physical_table( + "us", + "TIMESTAMP(6)", + ArrowTimeUnit::Microsecond, + [60_000_000, 120_000_000], + ) + .await; + run_otlp_logical_batcher_non_millisecond_physical_table( + "ns", + "TIMESTAMP(9)", + ArrowTimeUnit::Nanosecond, + [60_000_000_000, 120_000_000_000], + ) + .await; + } + + async fn run_otlp_logical_batcher_non_millisecond_physical_table( + suite: &str, + sql_ts_type: &str, + unit: ArrowTimeUnit, + expected: [i64; 2], + ) { + use std::time::Duration; + + use common_base::Plugins; + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::{TimestampMicrosecondType, TimestampNanosecondType}; + use frontend::server::Services; + use frontend::service_config::pending_rows_batcher::BatcherOptions; + use prost::Message; + use servers::batcher::BatchingProtocol; + use servers::http::test_helpers::TestClient; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + + let standalone = + GreptimeDbStandaloneBuilder::new(&format!("otlp_logical_{suite}_physical")) + .with_logical_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let options = standalone.opts.clone(); + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_logical_batching_enabled(true); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + + // Pre-create the physical metric table with a non-millisecond time + // index BEFORE any ingestion, so the batcher's bulk path must + // handle it. + let mut output = instance + .do_query( + &format!( + "CREATE TABLE greptime_physical_table (\ + greptime_timestamp {sql_ts_type} NOT NULL, \ + greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + ), + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + // submit_build_and_align counts every batcher submission in both + // acknowledgement modes, so the assertion detects a regression that + // disables non-millisecond batching. + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + let before = submissions.get_sample_count(); + for (ts, value) in [(60, 10), (120, 20)] { + let request = build_sum_request( + "non.ms.alignment", + AggregationTemporality::Cumulative, + &[(ts, value)], + ); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(response.status().as_u16(), 200); + } + assert_eq!( + submissions.get_sample_count() - before, + 2, + "non-millisecond physical tables must use the logical batcher" + ); + + // The rows land on the non-millisecond physical table, converted + // from the nanosecond encoding. + let sql = "SELECT greptime_timestamp, greptime_value FROM non_ms_alignment_total ORDER BY greptime_timestamp"; + let batches = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + if batches.iter().map(|batch| batch.num_rows()).sum::() == 2 { + break batches; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + let batch = &batches.take()[0]; + let stored = |row: usize| match unit { + ArrowTimeUnit::Microsecond => batch + .column(0) + .as_primitive::() + .value(row), + ArrowTimeUnit::Nanosecond => batch + .column(0) + .as_primitive::() + .value(row), + _ => unreachable!("covered units only"), + }; + assert_eq!(stored(0), expected[0]); + assert_eq!(stored(1), expected[1]); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap_or_else(|| panic!("expected f64 values")); + assert_eq!(values.value(0), 10.0); + assert_eq!(values.value(1), 20.0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_logical_batcher_fallback_for_cross_physical_destination() { + use std::time::Duration; + + use common_base::Plugins; + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::TimestampMicrosecondType; + use frontend::server::Services; + use frontend::service_config::pending_rows_batcher::BatcherOptions; + use prost::Message; + use servers::batcher::BatchingProtocol; + use servers::http::test_helpers::TestClient; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + + // With the logical batcher enabled, an OTLP request targeting an + // existing logical table bound to ANOTHER physical table must fall + // back to the ordinary insert path (the bulk eligibility check + // rejects the destination binding), which converts the request to + // the destination's unit. This covers the case previously guarded + // by the removed OTLP pre-gate alignment. + let standalone = GreptimeDbStandaloneBuilder::new("otlp_logical_cross_physical") + .with_logical_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let options = standalone.opts.clone(); + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_logical_batching_enabled(true); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + + // A physical table with a microsecond time index and a logical table + // bound to it; OTLP requests always select the default + // (millisecond) physical table, so the destination binding differs. + for sql in [ + "CREATE TABLE phy_us (\ + greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')", + "CREATE TABLE cross_alignment_total (\ + greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \ + \"stream\" STRING NULL, TIME INDEX (greptime_timestamp), PRIMARY KEY (\"stream\")) \ + ENGINE = metric WITH ('on_physical_table' = 'phy_us')", + ] { + let mut output = instance.do_query(sql, ctx.clone()).await; + let result = output.remove(0); + assert!(result.is_ok(), "setup ddl failed: {result:?} — {sql}"); + } + + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + let before = submissions.get_sample_count(); + let request = build_sum_request( + "cross.alignment", + AggregationTemporality::Cumulative, + &[(60, 10)], + ); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(response.status().as_u16(), 200); + // The request must NOT enter the logical batcher. + assert_eq!( + submissions.get_sample_count() - before, + 0, + "cross-physical destinations must fall back to the ordinary insert path" + ); + + // The row lands on the microsecond logical table, converted from the + // nanosecond encoding: 60s -> 60_000_000us. + let sql = "SELECT greptime_timestamp, greptime_value FROM cross_alignment_total"; + let batches = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + if batches.iter().map(|batch| batch.num_rows()).sum::() == 1 { + break batches; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + let batch = &batches.take()[0]; + let timestamps = batch.column(0).as_primitive::(); + assert_eq!(timestamps.value(0), 60_000_000); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_logical_batcher_alignment() { + use std::time::Duration; + + use common_base::Plugins; + use frontend::server::Services; + use frontend::service_config::pending_rows_batcher::BatcherOptions; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + ExponentialHistogram, ExponentialHistogramDataPoint, exponential_histogram_data_point, + }; + use prost::Message; + use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled}; + use servers::http::test_helpers::TestClient; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + + // Fixed workload, only the batcher opt-in changes. Compare stored rows, + // schema evolution and visibility under the configured acknowledgement policy. + let mut results = Vec::new(); + for enabled in [false, true] { + let standalone = GreptimeDbStandaloneBuilder::new(&format!("otlp_logical_{enabled}")) + .with_logical_batcher(BatcherOptions { + protocols: if enabled { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }, + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let options = standalone.opts.clone(); + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_logical_batching_enabled(true); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_wait_flush_result"]); + let before = submissions.get_sample_count(); + for (ts, value) in [(60, 10), (120, 20)] { + let mut request = build_sum_request( + "batch.alignment", + AggregationTemporality::Cumulative, + &[(ts, value)], + ); + if ts == 120 + && let Some(metric::Data::Sum(sum)) = + &mut request.resource_metrics[0].scope_metrics[0].metrics[0].data + { + sum.data_points[0].attributes.push(keyvalue("extra", "new")); + } + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(response.status().as_u16(), 200); + } + if enabled { + assert_eq!( + submissions.get_sample_count() - before, + if pending_rows_batch_sync_enabled() { + 2 + } else { + 0 + }, + "logical submissions must follow the global acknowledgement policy" + ); + } + let sql = "SELECT greptime_timestamp, greptime_value, stream, extra FROM batch_alignment_total ORDER BY greptime_timestamp"; + let batches = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + if !enabled + || pending_rows_batch_sync_enabled() + || batches.iter().map(|batch| batch.num_rows()).sum::() == 2 + { + break batches; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 2 + ); + results.push(batches.pretty_print().unwrap()); + + let before_histograms = submissions.get_sample_count(); + // First create a histogram table, then reuse it in a mixed export. + // Neither request may enter the scalar-only logical batcher. + for mixed in [false, true] { + let mut request = build_sum_request( + "batch.mixed", + AggregationTemporality::Cumulative, + &[(180, 30)], + ); + let metrics = &mut request.resource_metrics[0].scope_metrics[0].metrics; + let mut histogram = metrics[0].clone(); + histogram.name = "batch.histogram".to_string(); + histogram.data = Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + aggregation_temporality: AggregationTemporality::Cumulative as i32, + data_points: vec![ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_000, + time_unix_nano: if mixed { 4_000_000_000 } else { 3_000_000_000 }, + count: 4, + sum: Some(8.0), + zero_count: 1, + positive: Some(exponential_histogram_data_point::Buckets { + offset: -1, + bucket_counts: vec![1, 2], + }), + ..Default::default() + }], + })); + if !mixed { + metrics.clear(); + } + metrics.push(histogram); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!( + response.status().as_u16(), + 200, + "enabled={enabled}, mixed={mixed}" + ); + } + assert_eq!(submissions.get_sample_count(), before_histograms); + for (sql, expected_rows) in [ + ( + "SELECT greptime_timestamp, greptime_native_histogram FROM batch_histogram ORDER BY greptime_timestamp", + 2, + ), + ( + "SELECT greptime_timestamp, greptime_value FROM batch_mixed_total ORDER BY greptime_timestamp", + 1, + ), + ] { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + expected_rows + ); + results.push(batches.pretty_print().unwrap()); + } + + // An existing metric attached to another physical table must use + // the original routing, not the request's default physical table. + for sql in [ + "CREATE TABLE custom_physical (greptime_timestamp TIMESTAMP TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='')", + "CREATE TABLE custom_total (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value DOUBLE, \"stream\" STRING PRIMARY KEY) ENGINE=metric WITH ('on_physical_table'='custom_physical')", + ] { + instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + } + instance + .metrics( + build_sum_request("custom", AggregationTemporality::Cumulative, &[(60, 7)]), + ctx.clone(), + ) + .await + .unwrap(); + let output = instance + .do_query("SELECT greptime_value FROM custom_total", ctx.clone()) + .await + .remove(0) + .unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + assert!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap() + .contains("7.0") + ); + + // Request-level schema policy cannot be bypassed by batching. + let mut fixed = ctx.fork(); + fixed.set_extension("auto_create_table", "false"); + assert!( + instance + .metrics( + build_sum_request( + "missing", + AggregationTemporality::Cumulative, + &[(60, 1)] + ), + Arc::new(fixed) + ) + .await + .is_err() + ); + } + assert_eq!(results[..3], results[3..]); + } + #[tokio::test(flavor = "multi_thread")] pub async fn test_otlp_on_standalone() { let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_otlp") @@ -939,6 +2003,162 @@ WITH( ); } + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_microsecond_physical_table_on_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_us_physical") + .build() + .await; + + test_otlp_metrics_into_non_millisecond_physical_table( + standalone.fe_instance(), + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_704_067_200_123_456, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_microsecond_physical_table_on_distributed() { + let instance = tests::create_distributed_instance("test_otlp_us_physical_dist").await; + + test_otlp_metrics_into_non_millisecond_physical_table( + &instance.frontend(), + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_704_067_200_123_456, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_seconds_physical_table_on_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_s_physical") + .build() + .await; + + // Narrowing truncates: 1704067200123456789ns -> 1704067200s. + test_otlp_metrics_into_non_millisecond_physical_table( + standalone.fe_instance(), + "TIMESTAMP(0)", + TimeUnit::Second, + 1_704_067_200, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_nanoseconds_physical_table_on_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_ns_physical") + .build() + .await; + + // Same unit: the full nanosecond precision is kept verbatim. + test_otlp_metrics_into_non_millisecond_physical_table( + standalone.fe_instance(), + "TIMESTAMP(9)", + TimeUnit::Nanosecond, + 1_704_067_200_123_456_789, + ) + .await; + } + + /// Regression test for : + /// a physical metric table pre-created with a non-millisecond time index + /// accepts OTLP ingestion; the nanosecond samples are converted to the + /// physical table's unit (kept verbatim for a nanosecond table, + /// truncating the sub-unit part for micro/second tables). + async fn test_otlp_metrics_into_non_millisecond_physical_table( + instance: &Arc, + sql_ts_type: &str, + expected_unit: TimeUnit, + expected_value: i64, + ) { + let db = "otlp_non_ms_physical"; + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, db); + // Route the request to the metric engine, like the OTLP HTTP handler + // does with the default `prom_store.with_metric_engine = true`. + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + assert!( + SqlQueryHandler::do_query( + instance.as_ref(), + &format!("CREATE DATABASE IF NOT EXISTS {db}"), + ctx.clone(), + ) + .await + .first() + .unwrap() + .is_ok() + ); + + let mut output = instance + .do_query( + &format!( + "CREATE TABLE greptime_physical_table (\ + greptime_timestamp {sql_ts_type} NOT NULL, \ + greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + ), + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "my_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + attributes: vec![keyvalue("host", "h1")], + time_unix_nano: 1_704_067_200_123_456_789, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + instance.metrics(request, ctx.clone()).await.unwrap(); + + let mut output = instance + .do_query("SELECT greptime_timestamp FROM my_gauge", ctx.clone()) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = RecordBatches::try_collect(stream).await.unwrap().take(); + assert_eq!(batches[0].num_rows(), 1); + // The logical table's time index keeps the physical table's unit. + let ts_column = batches[0].column(0); + assert_eq!( + ts_column.data_type(), + &DataType::Timestamp(expected_unit, None), + "unexpected time index type" + ); + let stored = match expected_unit { + TimeUnit::Second => ts_column.as_primitive::().value(0), + TimeUnit::Millisecond => ts_column + .as_primitive::() + .value(0), + TimeUnit::Microsecond => ts_column + .as_primitive::() + .value(0), + TimeUnit::Nanosecond => ts_column.as_primitive::().value(0), + }; + assert_eq!(stored, expected_value); + } + fn build_request() -> ExportMetricsServiceRequest { let data_points = vec![ NumberDataPoint { diff --git a/tests-integration/src/prom_store.rs b/tests-integration/src/prom_store.rs index 0649b2dd778..af901120454 100644 --- a/tests-integration/src/prom_store.rs +++ b/tests-integration/src/prom_store.rs @@ -18,9 +18,19 @@ mod tests { use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{ - Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, WriteRequest, + Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, TimeSeries, WriteRequest, }; + use api::v1::value::ValueData; + use api::v1::{ + ColumnDataType, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, Value, + }; + use client::OutputData; use common_catalog::consts::DEFAULT_CATALOG_NAME; + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::{ + DataType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, + }; use frontend::instance::Instance; use prost::Message; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; @@ -83,6 +93,432 @@ mod tests { .await; } + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_remote_rw_microsecond_physical_table() { + common_telemetry::init_default_ut_logging(); + let standalone = + GreptimeDbStandaloneBuilder::new("test_prom_store_remote_rw_us_physical_table") + .build() + .await; + let instance = standalone.fe_instance(); + + test_prom_store_remote_rw_non_millisecond_physical_table( + instance, + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_000_000, + 2_000_000, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_distributed_prom_store_remote_rw_microsecond_physical_table() { + common_telemetry::init_default_ut_logging(); + let distributed = + tests::create_distributed_instance("test_prom_store_remote_rw_us_physical_table").await; + test_prom_store_remote_rw_non_millisecond_physical_table( + &distributed.frontend(), + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_000_000, + 2_000_000, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_remote_rw_seconds_physical_table() { + common_telemetry::init_default_ut_logging(); + let standalone = + GreptimeDbStandaloneBuilder::new("test_prom_store_remote_rw_s_physical_table") + .build() + .await; + let instance = standalone.fe_instance(); + + // Narrowing truncates: 1000ms/2000ms -> 1s/2s. + test_prom_store_remote_rw_non_millisecond_physical_table( + instance, + "TIMESTAMP(0)", + TimeUnit::Second, + 1, + 2, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_remote_rw_nanoseconds_physical_table() { + common_telemetry::init_default_ut_logging(); + let standalone = + GreptimeDbStandaloneBuilder::new("test_prom_store_remote_rw_ns_physical_table") + .build() + .await; + let instance = standalone.fe_instance(); + + // Lossless widening: 1000ms/2000ms -> 1e9ns/2e9ns. + test_prom_store_remote_rw_non_millisecond_physical_table( + instance, + "TIMESTAMP(9)", + TimeUnit::Nanosecond, + 1_000_000_000, + 2_000_000_000, + ) + .await; + } + + /// Regression test for : + /// prometheus remote write/read against a physical metric table + /// pre-created with a non-millisecond time index. Millisecond samples are + /// converted to the physical table's unit on write (lossless widening, + /// truncating narrowing) and narrowed back to milliseconds on remote read. + async fn test_prom_store_remote_rw_non_millisecond_physical_table( + instance: &Arc, + sql_ts_type: &str, + expected_unit: TimeUnit, + expected_first: i64, + expected_second: i64, + ) { + let db = "prometheus_non_ms"; + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, db)); + assert!( + SqlQueryHandler::do_query( + instance.as_ref(), + &format!("CREATE DATABASE IF NOT EXISTS {db}"), + ctx.clone(), + ) + .await + .first() + .unwrap() + .is_ok() + ); + + let mut output = instance + .do_query( + &format!( + "CREATE TABLE greptime_physical_table (\ + greptime_timestamp {sql_ts_type} NOT NULL, \ + greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + ), + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + let write_request = WriteRequest { + timeseries: vec![prom_store::mock_timeseries()[0].clone()], + ..Default::default() + }; + let (row_inserts, _) = to_grpc_row_insert_requests(&write_request).unwrap(); + instance + .write(row_inserts, ctx.clone(), true) + .await + .unwrap(); + + let read_request = ReadRequest { + queries: vec![Query { + start_timestamp_ms: 1000, + end_timestamp_ms: 2000, + matchers: vec![LabelMatcher { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "metric1".to_string(), + r#type: 0, + }], + ..Default::default() + }], + ..Default::default() + }; + let resp = instance.read(read_request, ctx.clone()).await.unwrap(); + let body = prom_store::snappy_decompress(&resp.body).unwrap(); + let read_response = ReadResponse::decode(&body[..]).unwrap(); + assert_eq!(1, read_response.results.len()); + assert_eq!(1, read_response.results[0].timeseries.len()); + let timeseries = &read_response.results[0].timeseries[0]; + assert_eq!( + timeseries.samples, + vec![ + Sample { + value: 1.0, + timestamp: 1000, + }, + Sample { + value: 2.0, + timestamp: 2000, + }, + ] + ); + + // The stored timestamps keep the physical table's time index unit. + let mut output = instance + .do_query("SELECT greptime_timestamp FROM metric1", ctx.clone()) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + .take(); + assert_eq!(batches[0].num_rows(), 2); + let ts_column = batches[0].column(0); + assert_eq!( + ts_column.data_type(), + &DataType::Timestamp(expected_unit, None), + "unexpected time index type" + ); + let stored = |row: usize| match expected_unit { + TimeUnit::Second => ts_column.as_primitive::().value(row), + TimeUnit::Millisecond => ts_column + .as_primitive::() + .value(row), + TimeUnit::Microsecond => ts_column + .as_primitive::() + .value(row), + TimeUnit::Nanosecond => ts_column + .as_primitive::() + .value(row), + }; + assert_eq!(stored(0), expected_first); + assert_eq!(stored(1), expected_second); + + // Negative, non-aligned timestamps must floor towards negative + // infinity on remote read, matching `Timestamp::convert_to` on the + // ingestion path: -1001us/-1001000ns -> -2ms, not -1ms. + if matches!(expected_unit, TimeUnit::Microsecond | TimeUnit::Nanosecond) { + let (ts_datatype, negative_value) = if expected_unit == TimeUnit::Microsecond { + ( + ColumnDataType::TimestampMicrosecond, + ValueData::TimestampMicrosecondValue(-1001), + ) + } else { + ( + ColumnDataType::TimestampNanosecond, + ValueData::TimestampNanosecondValue(-1_001_000), + ) + }; + let negative_request = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "metric1".to_string(), + rows: Some(Rows { + schema: vec![ + api::v1::ColumnSchema { + column_name: "greptime_timestamp".to_string(), + datatype: ts_datatype as i32, + semantic_type: SemanticType::Timestamp as i32, + datatype_extension: None, + options: None, + }, + api::v1::ColumnSchema { + column_name: "greptime_value".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension: None, + options: None, + }, + api::v1::ColumnSchema { + column_name: "job".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + datatype_extension: None, + options: None, + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(negative_value), + }, + Value { + value_data: Some(ValueData::F64Value(3.0)), + }, + Value { + value_data: Some(ValueData::StringValue("spark".to_string())), + }, + ], + }], + }), + }], + }; + instance + .write(negative_request, ctx.clone(), true) + .await + .unwrap(); + + let read_request = ReadRequest { + queries: vec![Query { + start_timestamp_ms: -2, + end_timestamp_ms: 0, + matchers: vec![LabelMatcher { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "metric1".to_string(), + r#type: 0, + }], + ..Default::default() + }], + ..Default::default() + }; + let resp = instance.read(read_request, ctx.clone()).await.unwrap(); + let body = prom_store::snappy_decompress(&resp.body).unwrap(); + let read_response = ReadResponse::decode(&body[..]).unwrap(); + assert_eq!(1, read_response.results.len()); + assert_eq!( + read_response.results[0].timeseries[0].samples, + vec![Sample { + value: 3.0, + timestamp: -2, + }], + "remote read must floor negative non-aligned timestamps" + ); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_write_existing_logical_table_of_other_physical_unit() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new( + "test_prom_store_write_existing_logical_other_physical_unit", + ) + .build() + .await; + let instance = standalone.fe_instance(); + + // Two physical metric tables with different time index units. + let db = "prometheus_mixed_units"; + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, db)); + assert!( + SqlQueryHandler::do_query( + instance.as_ref(), + &format!("CREATE DATABASE IF NOT EXISTS {db}"), + ctx.clone(), + ) + .await + .first() + .unwrap() + .is_ok() + ); + let mut output = instance + .do_query( + "CREATE TABLE phy_us (greptime_timestamp TIMESTAMP(6) NOT NULL, \ + greptime_value DOUBLE NULL, TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')", + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + // First create the logical table `shared_metric` on the default + // (millisecond) physical table. + let shared_series = TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "shared_metric".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }; + let (row_inserts, _) = to_grpc_row_insert_requests(&WriteRequest { + timeseries: vec![shared_series.clone()], + ..Default::default() + }) + .unwrap(); + instance + .write(row_inserts, ctx.clone(), true) + .await + .unwrap(); + + // Now select the microsecond physical table while writing to BOTH the + // existing millisecond `shared_metric` (bound to the default physical + // table) and a new table: the existing table's request must keep the + // millisecond unit, and only the new table uses microsecond. + let mut hint_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, db); + hint_ctx.set_extension(PHYSICAL_TABLE_PARAM, "phy_us".to_string()); + let hint_ctx = Arc::new(hint_ctx); + let (row_inserts, _) = to_grpc_row_insert_requests(&WriteRequest { + timeseries: vec![ + TimeSeries { + samples: vec![Sample { + value: 1.5, + timestamp: 1500, + }], + ..shared_series + }, + TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "fresh_us_metric".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { + value: 2.5, + timestamp: 1500, + }], + ..Default::default() + }, + ], + ..Default::default() + }) + .unwrap(); + instance + .write(row_inserts, hint_ctx.clone(), true) + .await + .unwrap(); + + // The existing table keeps its millisecond unit and precision. + let mut output = instance + .do_query( + "SELECT greptime_timestamp FROM shared_metric ORDER BY greptime_timestamp", + ctx.clone(), + ) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + .take(); + let shared_ts = batches[0] + .column(0) + .as_primitive::(); + assert_eq!((shared_ts.value(0), shared_ts.value(1)), (1000, 1500)); + + // The new table is created on the selected physical table with the + // microsecond unit. + let mut output = instance + .do_query( + "SELECT greptime_timestamp FROM fresh_us_metric", + ctx.clone(), + ) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + .take(); + let fresh_ts = batches[0] + .column(0) + .as_primitive::(); + assert_eq!(fresh_ts.value(0), 1_500_000); + } + async fn test_prom_store_remote_rw(instance: &Arc, physical_table: Option) { let write_request = WriteRequest { timeseries: prom_store::mock_timeseries(), diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 307a7f6daa2..241d4daf5b5 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use std::time::Duration; use cache::{ - build_datanode_cache_registry, build_fundamental_cache_registry, + build_datanode_layered_cache_registry, build_fundamental_cache_registry, with_default_composite_cache_registry, }; use catalog::information_schema::NoopInformationExtension; @@ -39,19 +39,19 @@ use common_meta::procedure_executor::LocalProcedureExecutor; use common_meta::region_keeper::MemoryRegionKeeper; use common_meta::region_registry::LeaderRegionRegistry; use common_meta::sequence::SequenceBuilder; -use common_meta::wal_provider::build_wal_provider; use common_procedure::ProcedureManagerRef; use common_procedure::local::EventRecorderHandle; use common_procedure::options::ProcedureConfig; use common_telemetry::logging::SlowQueryOptions; use common_test_util::find_workspace_path; -use common_wal::config::{DatanodeWalConfig, MetasrvWalConfig}; -use datanode::datanode::DatanodeBuilder; +use common_wal::config::DatanodeWalConfig; +use datanode::datanode::{Datanode, DatanodeBuilder}; use flow::{FlownodeBuilder, FrontendClient, GrpcQueryHandlerWithBoxedError}; use frontend::frontend::Frontend; use frontend::instance::Instance; use frontend::instance::builder::FrontendBuilder; use frontend::server::Services; +use frontend::service_config::{BatcherOptions, PendingRowsBatcherOptions}; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; use servers::grpc::GrpcOptions; use standalone::options::StandaloneOptions; @@ -62,6 +62,8 @@ use crate::test_util::{self, StorageType, TestGuard, create_tmp_dir_and_datanode pub struct GreptimeDbStandalone { /// Storage engine for assertions across protocol and storage boundaries. pub mito_engine: mito2::engine::MitoEngine, + /// Owns the region server and the log store, so a test can shut them down. + pub datanode: Datanode, pub frontend: Arc, pub opts: StandaloneOptions, pub guard: TestGuard, @@ -80,7 +82,6 @@ impl GreptimeDbStandalone { pub struct GreptimeDbStandaloneBuilder { instance_name: String, datanode_wal_config: DatanodeWalConfig, - metasrv_wal_config: MetasrvWalConfig, store_providers: Option>, default_store: Option, plugin: Option, @@ -88,6 +89,9 @@ pub struct GreptimeDbStandaloneBuilder { event_recorder_options: EventRecorderOptions, auto_create_table: bool, experimental_metric_export: bool, + logical_batcher: Option, + table_batcher: BatcherOptions, + mito_config: Option, } impl GreptimeDbStandaloneBuilder { @@ -98,7 +102,6 @@ impl GreptimeDbStandaloneBuilder { plugin: None, default_store: None, datanode_wal_config: DatanodeWalConfig::default(), - metasrv_wal_config: MetasrvWalConfig::default(), // Enable slow query log with 1s threshold by default for integration tests. slow_query_options: SlowQueryOptions { enable: true, @@ -108,9 +111,19 @@ impl GreptimeDbStandaloneBuilder { event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, experimental_metric_export: false, + logical_batcher: None, + table_batcher: BatcherOptions::default(), + mito_config: None, } } + /// Overrides the Mito configuration for this test instance. + #[must_use] + pub fn with_mito_config(mut self, config: mito2::config::MitoConfig) -> Self { + self.mito_config = Some(config); + self + } + /// Enables experimental Metric export for the standalone test instance. #[must_use] pub fn with_experimental_metric_export(mut self) -> Self { @@ -118,6 +131,20 @@ impl GreptimeDbStandaloneBuilder { self } + /// Configures ordinary-table batching for protocol integration tests. + #[must_use] + pub fn with_table_batcher(mut self, options: BatcherOptions) -> Self { + self.table_batcher = options; + self + } + + /// Configures logical-table batching for integration tests. + #[must_use] + pub fn with_logical_batcher(mut self, options: BatcherOptions) -> Self { + self.logical_batcher = Some(options); + self + } + #[must_use] pub fn with_auto_create_table(mut self, auto_create_table: bool) -> Self { self.auto_create_table = auto_create_table; @@ -171,12 +198,6 @@ impl GreptimeDbStandaloneBuilder { self } - #[must_use] - pub fn with_metasrv_wal_config(mut self, metasrv_wal_config: MetasrvWalConfig) -> Self { - self.metasrv_wal_config = metasrv_wal_config; - self - } - pub async fn build_with( &self, kv_backend: KvBackendRef, @@ -188,11 +209,8 @@ impl GreptimeDbStandaloneBuilder { ) -> GreptimeDbStandalone { let plugins = self.plugin.clone().unwrap_or_default(); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut builder = DatanodeBuilder::new(opts.datanode_options(), plugins.clone(), kv_backend.clone()); @@ -253,10 +271,13 @@ impl GreptimeDbStandaloneBuilder { .step(10) .build(), ); - let kafka_options = opts.wal.clone().try_into().unwrap(); - let wal_provider = build_wal_provider(&kafka_options, kv_backend.clone()) - .await - .unwrap(); + let wal_provider = cmd::standalone::build_standalone_wal_provider( + &opts.wal, + opts.datanode_options().node_id, + kv_backend.clone(), + ) + .await + .unwrap(); let wal_provider = Arc::new(wal_provider); let table_metadata_allocator = Arc::new(TableMetadataAllocator::new( table_id_allocator, @@ -335,6 +356,7 @@ impl GreptimeDbStandaloneBuilder { GreptimeDbStandalone { mito_engine: datanode.region_server().mito_engine().unwrap(), + datanode, frontend: Arc::new(frontend), opts, guard, @@ -360,23 +382,27 @@ impl GreptimeDbStandaloneBuilder { let procedure_config = ProcedureConfig::default(); let kv_backend = standalone::build_metadata_kvbackend( - format!("{}/kv", &opts.storage.data_home), + format!("{}/kv", opts.storage.data_home), kv_backend_config, ) .unwrap(); let (procedure_manager, event_recorder_handle) = standalone::build_procedure_manager(kv_backend.clone(), procedure_config); - let standalone_opts = StandaloneOptions { + let mut standalone_opts = StandaloneOptions { storage: opts.storage, procedure: procedure_config, metadata_store: kv_backend_config, - wal: self.metasrv_wal_config.clone().into(), + wal: opts.wal, grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, experimental_metric_export: self.experimental_metric_export, + pending_rows_batcher: PendingRowsBatcherOptions { + logical_table: self.logical_batcher.clone(), + table: self.table_batcher.clone(), + }, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { experimental_enable_resource_info: true, @@ -385,6 +411,14 @@ impl GreptimeDbStandaloneBuilder { ..StandaloneOptions::default() }; + if let Some(config) = &self.mito_config { + for engine in &mut standalone_opts.region_engine { + if let datanode::config::RegionEngineConfig::Mito(mito) = engine { + *mito = config.clone(); + } + } + } + self.build_with( kv_backend, guard, diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 562bd48cdb0..64f9f4c6af1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -31,6 +31,7 @@ use client::{Client, Database}; use common_base::Plugins; use common_catalog::consts::MIN_USER_TABLE_ID; use common_config::Configurable; +use common_event_recorder::EventRecorderOptions; #[cfg(test)] use common_meta::DatanodeId; use common_meta::key::TableMetadataManager; @@ -328,6 +329,18 @@ impl StorageType { } } +/// Event recorder options for tests. +/// +/// The production flush interval is 5s, and the event tests interleave "run a +/// DDL, wait for its event, run the next DDL", so each barrier costs a full +/// window. A short interval removes that wait without changing what is asserted. +pub fn test_event_recorder_options() -> EventRecorderOptions { + EventRecorderOptions { + flush_interval: Duration::from_millis(100), + ..Default::default() + } +} + fn s3_test_config() -> S3Config { S3Config { connection: S3Connection { @@ -363,7 +376,7 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Gcs::from(&gcs_config.connection); let config = ObjectStoreConfig::Gcs(gcs_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Gcs(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::Azblob => { let azblob_config = AzblobConfig { @@ -381,7 +394,7 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Azblob::from(&azblob_config.connection); let config = ObjectStoreConfig::Azblob(azblob_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Azblob(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::Oss => { let oss_config = OssConfig { @@ -398,32 +411,57 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Oss::from(&oss_config.connection); let config = ObjectStoreConfig::Oss(oss_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Oss(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::S3 | StorageType::S3WithCache => { let mut s3_config = s3_test_config(); - if *store_type == StorageType::S3WithCache { - s3_config.cache.cache_path = "/tmp/greptimedb_cache".to_string(); + // The datanode wipes `/cache/object/read` on startup, so a + // path shared between concurrently running tests lets a starting test + // delete the read cache of a running one. + let cache_dir = if *store_type == StorageType::S3WithCache { + let dir = create_temp_dir("gt_s3_read_cache"); + s3_config.cache.cache_path = dir.path().to_string_lossy().to_string(); + Some(dir) } else { s3_config.cache.enable_read_cache = false; - } + None + }; let builder = S3::from(&s3_config.connection); let config = ObjectStoreConfig::S3(s3_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::S3(TempFolder::new(&store, "/"))) + ( + config, + TempDirGuard { + remote: Some(TempFolder::new(&store, "/")), + local_cache: cache_dir, + }, + ) } - StorageType::File => (ObjectStoreConfig::File(FileConfig {}), TempDirGuard::None), + StorageType::File => ( + ObjectStoreConfig::File(FileConfig {}), + TempDirGuard::default(), + ), } } -pub enum TempDirGuard { - None, - S3(TempFolder), - Oss(TempFolder), - Azblob(TempFolder), - Gcs(TempFolder), +#[derive(Default)] +pub struct TempDirGuard { + /// Prefix to wipe from the remote object store, absent for the file backend. + remote: Option, + /// Local read cache directory. Only held so it is removed when the guard drops. + #[allow(dead_code)] + local_cache: Option, +} + +impl TempDirGuard { + fn remote(folder: TempFolder) -> Self { + Self { + remote: Some(folder), + local_cache: None, + } + } } pub struct TestGuard { @@ -446,11 +484,7 @@ pub struct StorageGuard(pub TempDirGuard); impl TestGuard { pub async fn remove_all(&mut self) { for storage_guard in self.storage_guards.iter_mut() { - if let TempDirGuard::S3(guard) - | TempDirGuard::Oss(guard) - | TempDirGuard::Azblob(guard) - | TempDirGuard::Gcs(guard) = &mut storage_guard.0 - { + if let Some(guard) = &mut storage_guard.0.remote { guard.remove_all().await.unwrap() } } @@ -465,11 +499,8 @@ impl Drop for TestGuard { common_runtime::spawn_global(async move { let mut errors = vec![]; for guard in guards { - if let TempDirGuard::S3(guard) - | TempDirGuard::Oss(guard) - | TempDirGuard::Azblob(guard) - | TempDirGuard::Gcs(guard) = guard.0 - && let Err(e) = guard.remove_all().await + if let Some(remote) = guard.0.remote + && let Err(e) = remote.remove_all().await { errors.push(e); } @@ -670,7 +701,7 @@ pub async fn setup_test_http_app_with_frontend_and_slow_query_threshold( .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler(instance.fe_instance().clone(), true, false) + .with_otlp_handler(instance.fe_instance().clone(), true) .with_jaeger_handler(instance.fe_instance().clone()) .with_greptime_config_options(instance.opts.to_toml().unwrap()) .build(); @@ -690,18 +721,6 @@ pub async fn setup_test_http_app_with_frontend_and_user_provider( user_provider, None, None, - false, - ) - .await -} - -pub async fn setup_test_http_app_with_otlp_exponential_histogram( - store_type: StorageType, - name: &str, - enabled: bool, -) -> (Router, TestGuard) { - setup_test_http_app_with_frontend_and_custom_options( - store_type, name, None, None, None, enabled, ) .await } @@ -712,7 +731,6 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( user_provider: Option, http_opts: Option, memory_limiter: Option, - experimental_enable_exponential_histogram: bool, ) -> (Router, TestGuard) { let plugins = Plugins::new(); if let Some(user_provider) = user_provider.clone() { @@ -729,16 +747,18 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( ..Default::default() }); - let mut http_server = HttpServerBuilder::new(http_opts) + // The memory limiter must be wired before the handlers: routes capture + // the builder's limiter at wiring time (mirrors `frontend::server`). + let mut http_server = HttpServerBuilder::new(http_opts); + if let Some(limiter) = memory_limiter { + http_server = http_server.with_memory_limiter(limiter); + } + let mut http_server = http_server .with_sql_handler(instance.fe_instance().clone()) .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler( - instance.fe_instance().clone(), - true, - experimental_enable_exponential_histogram, - ) + .with_otlp_handler(instance.fe_instance().clone(), true) .with_prometheus_handler(instance.fe_instance().clone()) .with_jaeger_handler(instance.fe_instance().clone()) .with_dashboard_handler(instance.fe_instance().clone()) @@ -748,10 +768,6 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( http_server = http_server.with_user_provider(user_provider); } - if let Some(limiter) = memory_limiter { - http_server = http_server.with_memory_limiter(limiter); - } - let http_server = http_server.build(); let app = http_server.build(http_server.make_app()).unwrap(); @@ -770,14 +786,7 @@ pub async fn setup_test_prom_app_with_frontend( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, false, false).await -} - -pub async fn setup_test_prom_app_with_frontend_native_histogram( - store_type: StorageType, - name: &str, -) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, false, true).await + setup_test_prom_app_with_frontend_inner(store_type, name, false, None).await } /// Like [`setup_test_prom_app_with_frontend`] but enables the pending-rows batcher, @@ -787,14 +796,25 @@ pub async fn setup_test_prom_app_with_frontend_batched( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, true, false).await + setup_test_prom_app_with_frontend_inner(store_type, name, true, None).await +} + +/// Like [`setup_test_prom_app_with_frontend`] but wires a shared request-memory +/// limiter into the server, mirroring production deployments that configure +/// `max_in_flight_write_bytes`. +pub async fn setup_test_prom_app_with_frontend_and_memory_limiter( + store_type: StorageType, + name: &str, + memory_limiter: Option, +) -> (Router, TestGuard) { + setup_test_prom_app_with_frontend_inner(store_type, name, false, memory_limiter).await } async fn setup_test_prom_app_with_frontend_inner( store_type: StorageType, name: &str, enable_batcher: bool, - experimental_enable_prometheus_native_histogram: bool, + memory_limiter: Option, ) -> (Router, TestGuard) { unsafe { std::env::set_var("TZ", "UTC"); @@ -846,7 +866,7 @@ async fn setup_test_prom_app_with_frontend_inner( let http_server = build_test_prom_server( instance.fe_instance().clone(), enable_batcher, - experimental_enable_prometheus_native_histogram, + memory_limiter, ) .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) .build(); @@ -855,10 +875,12 @@ async fn setup_test_prom_app_with_frontend_inner( } /// Builds Prometheus HTTP routes for either a standalone or distributed frontend. +/// `memory_limiter` must be set before the handlers are wired (handlers capture +/// the limiter at wiring time, mirroring production in `frontend::server`). pub fn build_test_prom_server( frontend_ref: Arc, enable_batcher: bool, - experimental_enable_prometheus_native_histogram: bool, + memory_limiter: Option, ) -> HttpServerBuilder { let http_opts = HttpOptions { addr: format!("127.0.0.1:{}", ports::get_port()), @@ -885,15 +907,18 @@ pub fn build_test_prom_server( None }; assert_eq!(pending_rows_batcher.is_some(), enable_batcher); - HttpServerBuilder::new(http_opts) + let mut builder = HttpServerBuilder::new(http_opts) .with_sql_handler(frontend_ref.clone()) - .with_logs_handler(frontend_ref.clone()) + .with_logs_handler(frontend_ref.clone()); + if let Some(limiter) = memory_limiter { + builder = builder.with_memory_limiter(limiter); + } + builder .with_prom_handler( frontend_ref.clone(), Some(frontend_ref.clone()), true, PromValidationMode::Strict, - experimental_enable_prometheus_native_histogram, pending_rows_batcher, ) .with_prometheus_handler(frontend_ref) @@ -977,8 +1002,14 @@ pub async fn setup_authenticated_grpc_database( .unwrap(); let grpc_server = setup_grpc_server_for_frontend_instance(instance, Some(user_provider)).await; let grpc_addr = grpc_server.bind_addr().unwrap().to_string(); - let mut database = - Database::new_with_dbname("greptime-public", Client::with_urls(vec![grpc_addr])); + let mut database = Database::new_with_dbname( + "greptime-public", + Client::with_query_and_control_managers( + Default::default(), + Default::default(), + vec![grpc_addr], + ), + ); database.set_auth(api::v1::auth_header::AuthScheme::Basic(Basic { username: username.to_string(), password: password.to_string(), @@ -1233,6 +1264,61 @@ pub async fn setup_pg_server_with_user_provider( (instance.guard, Arc::new(pg_server)) } +/// Sets up a standalone instance with both a Prometheus remote-write HTTP app +/// (native histograms enabled) and a Postgres server attached, so native +/// histogram data written via remote-write can be queried over the Postgres +/// protocol. +pub async fn setup_pg_server_with_prom_native_histogram( + store_type: StorageType, + name: &str, +) -> (TestGuard, Router, Arc>) { + unsafe { + std::env::set_var("TZ", "UTC"); + } + + let instance = setup_standalone_instance(name, store_type).await; + + // Prometheus remote-write HTTP app with native histograms enabled. + let http_server = build_test_prom_server(instance.fe_instance().clone(), false, None) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); + let app = http_server.build(http_server.make_app()).unwrap(); + + // Postgres server on the same instance. + let runtime = RuntimeBuilder::default() + .worker_threads(2) + .thread_name("pg-runtime") + .build() + .unwrap(); + + let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port()); + let opts = PostgresOptions { + addr: fe_pg_addr.clone(), + ..Default::default() + }; + let tls_server_config = Arc::new( + ReloadableTlsServerConfig::try_new(opts.tls.clone()) + .expect("Failed to load certificates and keys"), + ); + + let mut pg_server = Box::new(PostgresServer::new( + instance.fe_instance().clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + )); + + pg_server + .start(fe_pg_addr.parse::().unwrap()) + .await + .unwrap(); + + (instance.guard, app, Arc::new(pg_server)) +} + pub(crate) async fn prepare_another_catalog_and_schema(instance: &Instance) { let catalog_manager = instance .catalog_manager() diff --git a/tests-integration/src/tests.rs b/tests-integration/src/tests.rs index d70af021651..4a5b2f0c6c5 100644 --- a/tests-integration/src/tests.rs +++ b/tests-integration/src/tests.rs @@ -16,6 +16,7 @@ mod gc; mod instance_kafka_wal_test; mod instance_noop_wal_test; mod instance_test; +mod nested_merge_scan_capability_test; mod promql_test; mod reconcile_table; mod remote_dyn_filter_test; diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index dc4d3898860..a1ccac27328 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -2195,47 +2195,6 @@ async fn test_delete(instance: Arc) { check_output_stream(output, expect).await; } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_s3(instance: Arc) { - if let Ok(bucket) = env::var("GT_S3_BUCKET") - && !bucket.is_empty() - { - let instance = instance.frontend(); - - // setups - assert!(matches!(execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await.data, OutputData::AffectedRows(0))); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - let key_id = env::var("GT_S3_ACCESS_KEY_ID").unwrap(); - let key = env::var("GT_S3_ACCESS_KEY").unwrap(); - let region = env::var("GT_S3_REGION").unwrap(); - - let root = uuid::Uuid::new_v4().to_string(); - - // exports - let copy_to_stmt = format!( - "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - bucket, root, key_id, key, region - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_s3(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2268,9 +2227,21 @@ async fn test_execute_copy_from_s3(instance: Arc) { let key = env::var("GT_S3_ACCESS_KEY").unwrap(); let region = env::var("GT_S3_REGION").unwrap(); + let mut connection = format!( + "ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}'", + key_id, key, region + ); + // Honors an S3-compatible endpoint (MinIO in CI) so the test does not + // fall back to resolving the bucket against real AWS. + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") + && !endpoint.is_empty() + { + connection = format!("{},ENDPOINT='{}'", connection, endpoint); + } + let copy_to_stmt = format!( - "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - bucket, root, key_id, key, region + "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION ({})", + bucket, root, connection ); let output = execute_sql(&instance, ©_to_stmt).await.data; @@ -2315,10 +2286,7 @@ async fn test_execute_copy_from_s3(instance: Arc) { .data, OutputData::AffectedRows(0) )); - let sql = format!( - "{} CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - test.sql, key_id, key, region, - ); + let sql = format!("{} CONNECTION ({})", test.sql, connection); let output = execute_sql(&instance, &sql).await.data; assert!(matches!(output, OutputData::AffectedRows(2))); @@ -2340,50 +2308,6 @@ async fn test_execute_copy_from_s3(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_oss(instance: Arc) { - if let Ok(bucket) = env::var("GT_OSS_BUCKET") - && !bucket.is_empty() - { - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let key_id = env::var("GT_OSS_ACCESS_KEY_ID").unwrap(); - let key = env::var("GT_OSS_ACCESS_KEY").unwrap(); - let endpoint = env::var("GT_OSS_ENDPOINT").unwrap(); - - let root = uuid::Uuid::new_v4().to_string(); - - let copy_to_stmt = format!( - "Copy demo TO 'oss://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',ACCESS_KEY_SECRET='{}',ENDPOINT='{}')", - bucket, root, key_id, key, endpoint - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_oss(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2490,54 +2414,6 @@ async fn test_execute_copy_from_oss(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_gcs(instance: Arc) { - if let (Ok(bucket), Ok(scope), Ok(credential)) = ( - env::var("GT_GCS_BUCKET"), - env::var("GT_GCS_SCOPE"), - env::var("GT_GCS_CREDENTIAL"), - ) && !bucket.is_empty() - { - let endpoint = env::var("GT_GCS_ENDPOINT").unwrap_or_default(); - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let root = uuid::Uuid::new_v4().to_string(); - let mut connection = format!("SCOPE='{}',CREDENTIAL='{}'", scope, credential); - if !endpoint.is_empty() { - connection = format!("{},ENDPOINT='{}'", connection, endpoint); - } - - let copy_to_stmt = format!( - "Copy demo TO 'gcs://{}/{}/export/demo.parquet' CONNECTION ({})", - bucket, root, connection - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_gcs(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2646,50 +2522,6 @@ async fn test_execute_copy_from_gcs(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_azblob(instance: Arc) { - if let (Ok(container), Ok(account_name), Ok(account_key), Ok(endpoint)) = ( - env::var("GT_AZBLOB_CONTAINER"), - env::var("GT_AZBLOB_ACCOUNT_NAME"), - env::var("GT_AZBLOB_ACCOUNT_KEY"), - env::var("GT_AZBLOB_ENDPOINT"), - ) && !container.is_empty() - { - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let root = uuid::Uuid::new_v4().to_string(); - - let copy_to_stmt = format!( - "Copy demo TO 'azblob://{}/{}/export/demo.parquet' CONNECTION (ACCOUNT_NAME='{}',ACCOUNT_KEY='{}',ENDPOINT='{}')", - container, root, account_name, account_key, endpoint - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_azblob(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2816,7 +2648,7 @@ async fn test_execute_copy_from_orc_with_cast(instance: Arc) { let output = execute_sql( &instance, - &format!("copy demo from '{}' WITH(FORMAT='orc');", &filepath), + &format!("copy demo from '{}' WITH(FORMAT='orc');", filepath), ) .await .data; @@ -2855,7 +2687,7 @@ async fn test_execute_copy_from_orc(instance: Arc) { let output = execute_sql( &instance, - &format!("copy demo from '{}' WITH(FORMAT='orc');", &filepath), + &format!("copy demo from '{}' WITH(FORMAT='orc');", filepath), ) .await .data; @@ -2892,7 +2724,7 @@ async fn test_cast_type_issue_1594(instance: Arc) { let output = execute_sql( &instance, - &format!("copy tsbs_cpu from '{}' WITH(FORMAT='csv');", &filepath), + &format!("copy tsbs_cpu from '{}' WITH(FORMAT='csv');", filepath), ) .await; @@ -2930,7 +2762,7 @@ async fn test_copy_from_csv_skip_bad_records(instance: Arc) { &instance, &format!( "copy csv_skip_bad_records from '{}' WITH(FORMAT='csv', skip_bad_records='true');", - &filepath + filepath ), ) .await @@ -3480,3 +3312,166 @@ CREATE TABLE b ( +-------+-----------------------------------+"#; check_output_stream(output, expected).await; } + +#[rstest] +#[case::mito(false)] +#[case::batched_metric(true)] +#[tokio::test(flavor = "multi_thread")] +async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram, Sample}; + use axum::body::{Body, to_bytes}; + use http::Request; + use prost::Message; + use servers::http::{HttpOptions, HttpServerBuilder}; + use servers::prom_remote_write::v2::test_util as remote_write_v2; + use servers::prom_remote_write::validation::PromValidationMode; + use servers::prom_store::snappy_compress; + use tower::ServiceExt; + + use crate::cluster::GreptimeDbClusterBuilder; + use crate::test_util::build_test_prom_server; + use crate::tests::test_util::{MockInstanceBuilder, RebuildableMockInstance, TestContext}; + + common_telemetry::init_default_ut_logging(); + let builder = MockInstanceBuilder::Distributed( + GreptimeDbClusterBuilder::new(&format!("histogram_lifecycle_{metric_engine}")) + .await + .with_datanodes(3), + ); + let mut context = TestContext::new(builder).await; + let make_router = |frontend: Arc| { + let builder = if metric_engine { + build_test_prom_server(frontend.clone(), true, None) + } else { + HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(frontend.clone()) + .with_prometheus_handler(frontend.clone()) + .with_prom_handler( + frontend.clone(), + None, + false, + PromValidationMode::Strict, + None, + ) + }; + let server = builder.build(); + server.build(server.make_app()).unwrap() + }; + let snapshot = |frontend: Arc| async move { + let mut result = Vec::new(); + for table in ["lifecycle_sample", "lifecycle_histogram"] { + let output = execute_sql( + &frontend, + &format!("select * from {table} order by greptime_timestamp"), + ) + .await; + let OutputData::Stream(stream) = output.data else { + panic!("expected query stream") + }; + result.push( + util::collect_batches(stream) + .await + .unwrap() + .pretty_print() + .unwrap(), + ); + } + result + }; + let mut router = make_router(context.frontend()); + for round in 0..2 { + let timestamp = 1_000 + round * 10_000; + let mut labels = vec![("__name__", "lifecycle_histogram"), ("job", "api")]; + if round == 1 { + // Exercise schema evolution after the first SST already exists. + labels.push(("zone", "east")); + } + let sample = remote_write_v2::request_with_labels_and_samples( + vec![("__name__", "lifecycle_sample"), ("job", "api")], + vec![Sample { + value: 42.0, + timestamp, + ..Default::default() + }], + ); + let histograms = remote_write_v2::request_with_labels_and_histograms( + labels, + vec![Histogram { + timestamp, + start_timestamp: 500, + count: Some(Count::CountInt(9_007_199_254_740_993)), + zero_count: Some(ZeroCount::ZeroCountInt(1)), + zero_threshold: 0.001, + sum: 12.5, + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_deltas: vec![9_007_199_254_740_992], + ..Default::default() + }], + ); + // Create the scalar table first, then add histogram storage to its physical table. + for request in [sample, histograms] { + let response = router + .clone() + .oneshot( + Request::post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(Body::from( + snappy_compress(&request.encode_to_vec()).unwrap(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + response.status(), + 204, + "{:?}", + to_bytes(response.into_body(), usize::MAX).await.unwrap() + ); + } + + let expected = snapshot(context.frontend()).await; + assert!(expected[1].contains("9007199254740993"), "{}", expected[1]); + // Rebuild before flushing, including after adding fields to the physical + // schema: recovery must retain the columns filled into bulk writes. + drop(router); + context.rebuild().await; + router = make_router(context.frontend()); + assert_eq!(snapshot(context.frontend()).await, expected, "WAL recovery"); + let storage_tables = if metric_engine { + vec!["greptime_physical_table"] + } else { + vec!["lifecycle_sample", "lifecycle_histogram"] + }; + for table in &storage_tables { + execute_sql( + &context.frontend(), + &format!("admin flush_table('{table}')"), + ) + .await; + } + assert_eq!(snapshot(context.frontend()).await, expected, "SST read"); + if round == 1 { + for table in &storage_tables { + execute_sql( + &context.frontend(), + &format!("admin compact_table('{table}', 'strict_window', 'window=3600')"), + ) + .await; + } + assert_eq!(snapshot(context.frontend()).await, expected, "compaction"); + drop(router); + context.rebuild().await; + router = make_router(context.frontend()); + assert_eq!(snapshot(context.frontend()).await, expected, "SST reopen"); + } + } +} diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs new file mode 100644 index 00000000000..2c86589112d --- /dev/null +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -0,0 +1,749 @@ +// 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. + +//! Exercises hand-built nested `MergeScan` plans at datanode execution, not optimizer behavior. +//! The outer `handle_remote_read` request is in-process; inner region reads use real TCP Flight. +//! The probe scan resolves against the request's region catalog, while the nested payload resolves +//! against the engine catalog; distinct column names ensure these scopes cannot be confused. +//! Success compares the complete duplicate-sensitive result multiset with frontend SQL and known +//! rows. Missing inner tables and closed remote inner regions must each fail the whole query. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use api::v1::region::{QueryRequest, RegionRequestHeader}; +use client::Client; +use client::region::RegionRequester; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_grpc::channel_manager::ChannelManager; +use common_query::Output; +use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; +use common_telemetry::info; +use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; +use datanode::region_server::RegionServer; +use frontend::instance::Instance; +use query::datafusion::QUERY_PARALLELISM_HINT; +use query::dist_plan::MergeScanLogicalPlan; +use query::parser::QueryLanguageParser; +use query::query_engine::DefaultSerializer; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::{QueryContext, QueryContextRef}; +use store_api::region_request::{RegionCloseRequest, RegionRequest}; +use store_api::storage::RegionId; +use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; + +use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; +use crate::test_util::execute_sql; + +/// Partitioned probe table. +const PROBE_TABLE: &str = "nested_cap_probe"; +/// Partitioned build table with duplicate join keys. +const BUILD_TABLE: &str = "nested_cap_build"; +/// Missing inner table used to verify plan-resolution failure. +const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; + +/// `(a_id, probe_key, probe_v)` fixture rows; `a_id` is the partition key. +const PROBE_ROWS: &[(i32, i32, i32)] = &[ + (1, 10, 100), + (2, 20, 200), + // The second partition of the probe table. + (120, 30, 300), + (121, 10, 400), +]; + +/// `(b_id, build_key, build_v)` fixture rows; key `10` duplicates and `40` is unmatched. +const BUILD_ROWS: &[(i32, i32, i32)] = &[ + (1, 10, 1), + (2, 10, 2), + (3, 20, 3), + // The second partition of the build table. + (120, 30, 4), + (121, 40, 5), +]; + +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_join_across_datanodes() { + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster("test_nested_merge_scan_capability_join_across_datanodes").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + assert_eq!( + 2, + datanodes(&probe_leaders).len(), + "expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \ + leaders: {probe_leaders:?}" + ); + assert_eq!( + 2, + datanodes(&build_leaders).len(), + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \ + leaders: {build_leaders:?}" + ); + + let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; + assert_eq!( + expected_join_rows(), + multiset(table_cells(&reference)), + "unexpected result of the join on the frontend:\n{reference}" + ); + + let query_ctx = query_ctx(); + let nested_plan = nested_plan(&frontend, &query_ctx).await; + let plan = encode(&nested_plan); + + let remote_datanodes = datanodes(&build_leaders); + let baseline_requests = remote_datanodes + .iter() + .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) + .collect::>(); + + let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; + + let actual_rows = multiset(rows); + assert_eq!( + expected_join_rows(), + actual_rows, + "the nested plans of the datanodes returned an unexpected multiset of rows" + ); + assert_eq!( + multiset(table_cells(&reference)), + actual_rows, + "frontend result differs:\n{reference}" + ); + assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); +} + +/// Internal read-preference wire metadata is rejected before nested execution on both inbound paths. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_rejects_internal_read_preference() { + use common_meta::node_manager::Datanode; + use common_query::request::QueryRequest as LogicalQueryRequest; + use futures::TryStreamExt; + use session::hints::READ_PREFERENCE_EXTENSION_KEY; + + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster("test_nested_merge_scan_internal_read_preference").await; + let frontend = cluster.fe_instance().clone(); + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let query_ctx = query_ctx(); + let plan = nested_plan(&frontend, &query_ctx).await; + let (region_id, (datanode_id, address)) = probe_leaders.iter().next().unwrap(); + let target = region_server(&cluster, *datanode_id); + // Confirm the same plan completes on the valid default-Leader path before asserting rejection. + let valid_rows = run_nested_queries(&cluster, &probe_leaders, &encode(&plan), &query_ctx).await; + assert_eq!(expected_join_rows(), multiset(valid_rows)); + let mut baseline = cluster + .datanode_instances + .keys() + .map(|id| (*id, rpc_requests(&cluster, *id))) + .collect::>(); + + for value in ["", "invalid", "follower"] { + let mut header = region_request_header(&query_ctx); + header + .query_context + .as_mut() + .unwrap() + .extensions + .insert(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string()); + let header = Some(header); + let error = target + .handle_read(LogicalQueryRequest { + header: header.clone(), + region_id: RegionId::from_u64(*region_id), + plan: plan.clone(), + }) + .await + .err() + .expect("internal read preference must be rejected at logical ingress"); + assert_eq!( + StatusCode::InvalidArguments, + error.status_code(), + "value={value:?}" + ); + let logical_message = error.output_msg().to_ascii_lowercase(); + assert!( + logical_message.contains("read_preference"), + "expected an invalid query-context read-preference error, got: {logical_message}" + ); + if !value.is_empty() { + assert!( + logical_message.contains(value), + "unexpected logical error: {logical_message}" + ); + } + for id in cluster.datanode_instances.keys() { + assert_eq!( + baseline[id], + rpc_requests(&cluster, *id), + "logical invalid context caused DoGet at datanode {id}" + ); + } + + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [address.as_str()], + ); + let requester = RegionRequester::new(client, false, false); + let request = LogicalQueryRequest { + header, + region_id: RegionId::from_u64(*region_id), + plan: plan.clone(), + }; + let flight_message = match Datanode::handle_query(&requester, request).await { + Err(error) => { + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + error.output_msg().to_ascii_lowercase() + } + Ok(stream) => { + let error = stream + .try_collect::>() + .await + .expect_err("Flight must reject invalid internal read preference"); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + error.output_msg().to_ascii_lowercase() + } + }; + assert!( + flight_message.contains("read_preference"), + "expected an invalid-parameter read-preference error, got: {flight_message}" + ); + if !value.is_empty() { + assert!( + flight_message.contains(value), + "unexpected Flight error: {flight_message}" + ); + } + for id in cluster.datanode_instances.keys() { + let expected = baseline[id] + u8::from(*id == *datanode_id) as usize; + assert_eq!( + expected, + rpc_requests(&cluster, *id), + "invalid Flight context caused nested DoGet at datanode {id}" + ); + } + for id in cluster.datanode_instances.keys() { + baseline.insert(*id, rpc_requests(&cluster, *id)); + } + } +} + +/// A failing inner region makes the whole query fail: the datanode must not return the rows of the +/// inner regions that succeed. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_inner_failure_fails_query() { + common_telemetry::init_default_ut_logging(); + + let cluster = + build_cluster("test_nested_merge_scan_capability_inner_failure_fails_query").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + + let (probe_region, (probe_datanode, _)) = probe_leaders + .iter() + .find(|(_, (datanode, _))| *datanode == 1) + .unwrap_or_else(|| { + panic!( + "expected a region of {PROBE_TABLE} on the datanode 1, actual region leaders: \ + {probe_leaders:?}" + ) + }); + let (build_region, (build_datanode, _)) = build_leaders + .iter() + .find(|(_, (datanode, _))| *datanode != *probe_datanode) + .unwrap_or_else(|| { + panic!( + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual \ + region leaders: {build_leaders:?}" + ) + }); + + // The nested payload resolves against the engine catalog, not the request's region catalog. + let query_ctx = query_ctx(); + let missing_inner_plan = + nested_plan_with_inner_table(&frontend, &query_ctx, MISSING_INNER_TABLE).await; + let plan = encode(&missing_inner_plan); + let error = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan, + }, + query_ctx.clone(), + ) + .await + .err() + .expect("a nested plan whose inner relation does not exist must not be executed"); + // The top level error of the datanode only names the failing step (`Failed to decode logical + // plan`): the reason is the source of it, which names the table that could not be resolved. + let chain = error_chain(&error); + assert!( + chain.contains(MISSING_INNER_TABLE), + "expected the datanode to fail while resolving the inner relation {MISSING_INNER_TABLE}, \ + actual error: {chain}" + ); + + // A closed remote inner region must fail the whole query, not return a partial result. + region_server(&cluster, *build_datanode) + .handle_request( + RegionId::from_u64(*build_region), + RegionRequest::Close(RegionCloseRequest { + flush_on_close: false, + }), + ) + .await + .expect("the datanode must close the region of the build table"); + + let nested_plan = nested_plan(&frontend, &query_ctx).await; + let plan = encode(&nested_plan); + let baseline = rpc_requests(&cluster, *build_datanode); + let stream = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan, + }, + query_ctx.clone(), + ) + .await + .expect("the request itself must be accepted by the datanode"); + let error = collect_error(stream) + .await + .expect("the query must fail when one of its inner regions fails"); + info!("the nested plan failed as a whole: {error}"); + + // The increased DoGet count proves execution reached the remote region owner. + let actual = rpc_requests(&cluster, *build_datanode); + assert!( + actual > baseline, + "expected the failed nested plan to reach datanode {build_datanode} for region {build_region} over Flight DoGet (count {actual}, baseline {baseline})" + ); +} + +async fn build_cluster(test_name: &str) -> GreptimeDbCluster { + GreptimeDbClusterBuilder::new(test_name) + .await + // A datanode serves the inner `MergeScan` of a nested plan by querying the datanodes that + // own the regions of the inner table, so every datanode must be reachable at the address it + // registers in metasrv, i.e. the address of the table routes. The cluster then also counts + // the requests of the remote nodes, see [`assert_cross_datanode_region_query`]. + .with_real_datanode_grpc_addr(true) + .with_datanodes(2) + .build(false) + .await +} + +fn assert_cross_datanode_region_query( + cluster: &GreptimeDbCluster, + remote_datanodes: &BTreeSet, + baseline_requests: &BTreeMap, +) { + assert_eq!( + 2, + remote_datanodes.len(), + "build table must span both datanodes" + ); + for datanode in remote_datanodes { + let actual = rpc_requests(cluster, *datanode); + let baseline = baseline_requests[datanode]; + assert!( + actual > baseline, + "expected a nested Flight DoGet at datanode {datanode}, count {actual}, baseline {baseline}" + ); + } +} + +fn rpc_requests(cluster: &GreptimeDbCluster, datanode_id: u64) -> usize { + cluster + .datanode_rpc_stats(datanode_id) + .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode_id}")) + .requests() +} + +fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer { + cluster + .datanode_instances + .get(&datanode_id) + .unwrap_or_else(|| panic!("expected a datanode {datanode_id}")) + .region_server() +} + +async fn run_nested_queries( + cluster: &GreptimeDbCluster, + probe_leaders: &BTreeMap, + plan: &[u8], + query_ctx: &QueryContextRef, +) -> Vec> { + let mut rows = Vec::new(); + for (probe_region, (probe_datanode, _)) in probe_leaders { + let stream = region_server(cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(query_ctx)), + region_id: *probe_region, + plan: plan.to_vec(), + }, + query_ctx.clone(), + ) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region}: {e}" + ) + }); + let batches = RecordBatches::try_collect(stream) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region} completely: {e}" + ) + }); + let actual = batches + .pretty_print() + .expect("the result of the datanode must be printable"); + rows.extend(table_cells(&actual)); + } + + rows +} + +/// Creates the probe and build tables of the tests and inserts the rows. +async fn prepare_tables(frontend: &Arc) { + execute_sql( + frontend, + &format!( + "CREATE TABLE {PROBE_TABLE} ( + a_id INT, + probe_key INT, + probe_v INT, + a_ts TIMESTAMP, + TIME INDEX (a_ts), + PRIMARY KEY(a_id) + ) + PARTITION ON COLUMNS (a_id) ( + a_id < 100, + a_id >= 100 + ) + engine=mito" + ), + ) + .await; + execute_sql( + frontend, + &format!( + "CREATE TABLE {BUILD_TABLE} ( + b_id INT, + build_key INT, + build_v INT, + b_ts TIMESTAMP, + TIME INDEX (b_ts), + PRIMARY KEY(b_id) + ) + PARTITION ON COLUMNS (b_id) ( + b_id < 100, + b_id >= 100 + ) + engine=mito" + ), + ) + .await; + + execute_sql( + frontend, + &format!( + "INSERT INTO {PROBE_TABLE}(a_id, probe_key, probe_v, a_ts) VALUES {}", + PROBE_ROWS + .iter() + .enumerate() + .map(|(i, (a_id, probe_key, probe_v))| { + format!("({a_id}, {probe_key}, {probe_v}, {}000)", i + 1) + }) + .collect::>() + .join(",") + ), + ) + .await; + execute_sql( + frontend, + &format!( + "INSERT INTO {BUILD_TABLE}(b_id, build_key, build_v, b_ts) VALUES {}", + BUILD_ROWS + .iter() + .enumerate() + .map(|(i, (b_id, build_key, build_v))| { + format!("({b_id}, {build_key}, {build_v}, {}000)", i + 1) + }) + .collect::>() + .join(",") + ), + ) + .await; +} + +fn encode(plan: &LogicalPlan) -> Vec { + DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .expect("the plan must be encodable") + .to_vec() +} + +/// Plans the hand-built nested join sent in the region request. +async fn nested_plan(frontend: &Arc, query_ctx: &QueryContextRef) -> LogicalPlan { + nested_plan_with_inner_table(frontend, query_ctx, BUILD_TABLE).await +} + +async fn nested_plan_with_inner_table( + frontend: &Arc, + query_ctx: &QueryContextRef, + inner_table: &str, +) -> LogicalPlan { + let probe = plan_sql( + frontend, + &format!("SELECT a_id, probe_key, probe_v FROM {PROBE_TABLE}"), + query_ctx, + ) + .await; + let build = plan_sql( + frontend, + &format!("SELECT build_key, build_v FROM {BUILD_TABLE}"), + query_ctx, + ) + .await; + let build = rename_scans(build, inner_table); + + nested_plan_of(&probe, build) +} + +fn nested_plan_of(probe: &LogicalPlan, build: LogicalPlan) -> LogicalPlan { + info!("probe side of the nested plan:\n{probe}"); + info!("build side of the nested plan:\n{build}"); + + let inner = MergeScanLogicalPlan::new(build, false, Default::default()).into_logical_plan(); + let plan = LogicalPlanBuilder::from(probe.clone()) + .join( + inner, + JoinType::Inner, + (vec!["probe_key"], vec!["build_key"]), + None, + ) + .expect("the join keys must be resolvable in the sides of the join") + .build() + .expect("the nested plan must be a valid logical plan"); + info!("nested plan sent to the datanode:\n{plan}"); + + plan +} + +fn rename_scans(plan: LogicalPlan, table_name: &str) -> LogicalPlan { + use datafusion::common::TableReference; + use datafusion::common::tree_node::{Transformed, TreeNode}; + + plan.transform_up(|plan| match plan { + LogicalPlan::TableScan(mut scan) => { + info!( + "renaming the scan of the nested plan from {} to {table_name}", + scan.table_name + ); + scan.table_name = TableReference::bare(table_name); + Ok(Transformed::yes(LogicalPlan::TableScan(scan))) + } + other => Ok(Transformed::no(other)), + }) + .expect("the plan must be rewritable") + .data +} + +async fn plan_sql(frontend: &Arc, sql: &str, query_ctx: &QueryContextRef) -> LogicalPlan { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx) + .unwrap_or_else(|e| panic!("failed to parse `{sql}`: {e}")); + frontend + .statement_executor() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap_or_else(|e| panic!("failed to plan `{sql}`: {e}")) +} + +/// Frontend reference join over the nested plan's output columns. +fn join_sql() -> String { + format!( + "SELECT p.a_id, p.probe_key, p.probe_v, b.build_key, b.build_v + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b ON p.probe_key = b.build_key" + ) +} + +/// Expected complete join result, retaining duplicate rows. +fn expected_join_rows() -> Vec> { + let mut rows = Vec::new(); + for (a_id, probe_key, probe_v) in PROBE_ROWS { + for (_, build_key, build_v) in BUILD_ROWS { + if probe_key == build_key { + rows.push(vec![ + a_id.to_string(), + probe_key.to_string(), + probe_v.to_string(), + build_key.to_string(), + build_v.to_string(), + ]); + } + } + } + multiset(rows) +} + +async fn region_leaders(frontend: &Arc, table: &str) -> BTreeMap { + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::UInt64Type; + + let sql = format!( + "SELECT region_id, peer_id, peer_addr FROM information_schema.region_peers + WHERE table_schema = 'public' AND table_name = '{table}' AND is_leader = 'Yes' + ORDER BY region_id" + ); + let batches = output_batches(run_sql(frontend, &sql, query_ctx()).await).await; + + let mut leaders = BTreeMap::new(); + for batch in batches.iter() { + let batch = batch.df_record_batch(); + let region_ids = batch + .column_by_name("region_id") + .expect("region_id column") + .as_primitive::(); + let peer_ids = batch + .column_by_name("peer_id") + .expect("peer_id column") + .as_primitive::(); + let peer_addrs = batch + .column_by_name("peer_addr") + .expect("peer_addr column") + .as_string::(); + for row in 0..batch.num_rows() { + leaders.insert( + region_ids.value(row), + (peer_ids.value(row), peer_addrs.value(row).to_string()), + ); + } + } + + leaders +} + +fn datanodes(leaders: &BTreeMap) -> BTreeSet { + leaders.values().map(|(id, _)| *id).collect() +} + +async fn output_batches(output: Output) -> RecordBatches { + match output.data { + common_query::OutputData::Stream(stream) => RecordBatches::try_collect(stream) + .await + .expect("the output stream must be collectable"), + common_query::OutputData::RecordBatches(batches) => batches, + common_query::OutputData::AffectedRows(rows) => { + panic!("expected a query output, got {rows} affected rows") + } + } +} + +/// Consumes the full stream to detect errors after any partial rows. +async fn collect_error(stream: SendableRecordBatchStream) -> Option { + use futures::TryStreamExt; + + let result = stream.try_collect::>().await; + match result { + Ok(rows) => { + info!( + "the query returned {} batches instead of failing", + rows.len() + ); + None + } + Err(error) => Some(error_chain(&error)), + } +} + +fn error_chain(error: &(dyn std::error::Error + 'static)) -> String { + let mut messages = vec![error.to_string()]; + let mut source = error.source(); + while let Some(inner) = source { + messages.push(inner.to_string()); + source = inner.source(); + } + messages.join(": ") +} + +async fn query_pretty(frontend: &Arc, sql: &str, query_ctx: QueryContextRef) -> String { + run_sql(frontend, sql, query_ctx) + .await + .data + .pretty_print() + .await +} + +async fn run_sql(frontend: &Arc, sql: &str, query_ctx: QueryContextRef) -> Output { + SqlQueryHandler::do_query(frontend.as_ref(), sql, query_ctx) + .await + .remove(0) + .unwrap() +} + +/// Uses one target partition to keep the plans stable. +fn query_ctx() -> QueryContextRef { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, "1"); + Arc::new(query_ctx) +} + +fn region_request_header(query_ctx: &QueryContextRef) -> RegionRequestHeader { + RegionRequestHeader { + query_context: Some(query_ctx.as_ref().into()), + ..Default::default() + } +} + +fn table_cells(pretty: &str) -> Vec> { + let mut lines = pretty.lines().filter(|line| line.starts_with('|')); + // The first `|` line is the header of the table. + lines.next(); + + lines + .map(|line| { + line.trim_matches('|') + .split('|') + .map(|cell| cell.trim().to_string()) + .collect() + }) + .collect() +} + +fn multiset(mut rows: Vec>) -> Vec> { + rows.sort(); + rows +} diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index c816583f135..74114bbd6c4 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -644,28 +644,57 @@ async fn two_aggregators_combined_aggrs(instance: Arc) { // eval instant at 50m stddev by (instance)(http_requests) // {instance="0"} 223.60679774998 // {instance="1"} 223.60679774998 +// +// The values are compared with a tolerance rather than against a rendered table: +// `stddev_pop` merges partial aggregates in a parallelism-dependent order, so the +// last digits of its output are not stable across runs. #[apply(both_instances_cases)] -#[ignore = "TODO(ruihang): fix this case"] async fn stddev_by_label(instance: Arc) { let instance = instance.frontend(); - create_insert_query_assert( + execute_all(&instance, AGGREGATORS_CREATE_TABLE, QueryContext::arc()).await; + execute_all(&instance, AGGREGATORS_INSERT_DATA, QueryContext::arc()).await; + + let batches = promql_query_as_batches( instance, - AGGREGATORS_CREATE_TABLE, - AGGREGATORS_INSERT_DATA, r#"stddev by (instance)(http_requests)"#, + None, + QueryContext::arc(), UNIX_EPOCH, unix_epoch_plus_100s(), Duration::from_secs(60), Duration::from_secs(0), - "+----------+---------------------+--------------------------------+\ - \n| instance | ts | STDDEVPOP(http_requests.value) |\ - \n+----------+---------------------+--------------------------------+\ - \n| 0 | 1970-01-01T00:00:00 | 223.606797749979 |\ - \n| 1 | 1970-01-01T00:00:00 | 223.606797749979 |\ - \n+----------+---------------------+--------------------------------+", ) .await; + + let batches = batches.iter().collect::>(); + assert_eq!(batches.len(), 1); + // Tag columns come back dictionary-encoded. + let instances = cast( + batches[0].column_by_name("instance").unwrap(), + &DataType::Utf8, + ) + .unwrap(); + let instances = instances.as_any().downcast_ref::().unwrap(); + let values = batches[0] + .column_by_name("stddev_pop(http_requests.value)") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + let mut actual = (0..values.len()) + .map(|i| (instances.value(i), values.value(i))) + .collect::>(); + actual.sort_by(|a, b| a.0.cmp(b.0)); + assert_eq!(actual.len(), 2, "{}", batches[0].pretty_print()); + for (expected_instance, (instance, value)) in ["0", "1"].into_iter().zip(actual) { + assert_eq!(instance, expected_instance); + assert!( + (value - 223.606_797_749_979).abs() < 1e-9, + "instance {instance}: {value}" + ); + } } // This is not derived from prometheus diff --git a/tests-integration/src/tests/test_util.rs b/tests-integration/src/tests/test_util.rs index 5ddcd0748a9..7c99e89ee08 100644 --- a/tests-integration/src/tests/test_util.rs +++ b/tests-integration/src/tests/test_util.rs @@ -263,15 +263,8 @@ pub(crate) async fn standalone_with_kafka_wal() -> Option>(); let test_name = uuid::Uuid::new_v4().to_string(); - let builder = GreptimeDbStandaloneBuilder::new(&test_name) - .with_datanode_wal_config(DatanodeWalConfig::Kafka(DatanodeKafkaConfig { - connection: KafkaConnectionConfig { - broker_endpoints: endpoints.clone(), - ..Default::default() - }, - ..Default::default() - })) - .with_metasrv_wal_config(MetasrvWalConfig::Kafka(MetasrvKafkaConfig { + let builder = GreptimeDbStandaloneBuilder::new(&test_name).with_datanode_wal_config( + DatanodeWalConfig::Kafka(DatanodeKafkaConfig { connection: KafkaConnectionConfig { broker_endpoints: endpoints, ..Default::default() @@ -282,7 +275,8 @@ pub(crate) async fn standalone_with_kafka_wal() -> Option) { execute( instance, diff --git a/tests-integration/tests/build_series_index.rs b/tests-integration/tests/build_series_index.rs new file mode 100644 index 00000000000..8c0e91aa35b --- /dev/null +++ b/tests-integration/tests/build_series_index.rs @@ -0,0 +1,148 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +use datanode::config::RegionEngineConfig; +use frontend::instance::Instance; +use mito2::config::MitoConfig; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::{ + StorageType, create_tmp_dir_and_datanode_opts, execute_sql, try_execute_sql, +}; + +fn config() -> MitoConfig { + MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + } +} + +async fn exercise_admin(frontend: &Arc) { + execute_sql( + frontend, + r#" + CREATE TABLE series_physical (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) + PARTITION ON COLUMNS (host) (host < 'm', host >= 'm') + ENGINE = metric WITH ( + physical_metric_table = 'true', 'compaction.type' = 'twcs', + 'compaction.twcs.time_window' = '1h', + 'compaction.twcs.active_window.trigger_file_num' = '100', + 'compaction.twcs.inactive_window.trigger_file_num' = '100' + ) + "#, + ) + .await; + execute_sql( + frontend, + r#" + CREATE TABLE series_logical (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) + ENGINE = metric WITH (on_physical_table = 'series_physical') + "#, + ) + .await; + // Reconciliation requires at least four SSTs in each bucket. + for timestamp in [1000, 2000, 3000, 4000] { + execute_sql( + frontend, + &format!( + "INSERT INTO series_logical (ts, host, val) VALUES ({timestamp}, 'a', 1), ({timestamp}, 'z', 2)" + ), + ) + .await; + execute_sql(frontend, "ADMIN FLUSH_TABLE('series_physical')").await; + } + for _ in 0..2 { + let output = execute_sql(frontend, "ADMIN BUILD_SERIES_INDEX('series_physical')").await; + assert!(output.data.pretty_print().await.contains("| 0")); + } + // The function must never silently expand a logical table to its shared physical table. + let error = try_execute_sql(frontend, "ADMIN BUILD_SERIES_INDEX('series_logical')") + .await + .unwrap_err(); + assert!(format!("{error:?}").contains("physical metric table")); + for sql in [ + "ADMIN BUILD_SERIES_INDEX()", + "ADMIN BUILD_SERIES_INDEX(1)", + "ADMIN BUILD_SERIES_INDEX('missing')", + ] { + assert!(try_execute_sql(frontend, sql).await.is_err(), "{sql}"); + } + execute_sql(frontend, "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX)").await; + assert!( + try_execute_sql(frontend, "ADMIN BUILD_SERIES_INDEX('ordinary')") + .await + .is_err() + ); + execute_sql(frontend, "ADMIN BUILD_INDEX('ordinary')").await; + let output = execute_sql(frontend, "SELECT count(*) AS n FROM series_logical").await; + assert!(output.data.pretty_print().await.contains("| 8")); +} + +fn catalog_count(home: &str) -> usize { + let root = Path::new(home).join("series_index"); + std::fs::read_dir(root) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| path.join("series-index.json").is_file()) + .count() +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_build_series_index_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("build-series-index") + .with_mito_config(config()) + .build() + .await; + exercise_admin(standalone.fe_instance()).await; + assert_eq!(2, catalog_count(&standalone.opts.storage.data_home)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_build_series_index_distributed() { + let builder = GreptimeDbClusterBuilder::new("build-series-index").await; + let mut options = Vec::new(); + let mut guards = Vec::new(); + for id in 1..=2 { + let (mut opts, guard) = create_tmp_dir_and_datanode_opts( + StorageType::File, + vec![], + &format!("build-series-index-{id}"), + Default::default(), + Default::default(), + ); + opts.node_id = Some(id); + for engine in &mut opts.region_engine { + if let RegionEngineConfig::Mito(mito) = engine { + *mito = config(); + } + } + options.push(opts); + guards.push(guard); + } + let cluster = builder.build_with(options, false, guards).await; + exercise_admin(cluster.fe_instance()).await; + assert_eq!( + 2, + cluster + .datanode_options + .iter() + .map(|opts| catalog_count(&opts.storage.data_home)) + .sum::() + ); +} diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index c08b81211de..1c6d3f39633 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -21,6 +21,7 @@ use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ StorageType, execute_sql, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use crate::event_recorder_test_util::{ @@ -42,7 +43,7 @@ async fn test_event_table_auto_creation_with_auto_create_disabled() { event_types: Arc::new(EventTypeFilter::Only( [String::from("create_database")].into_iter().collect(), )), - ..Default::default() + ..test_event_recorder_options() }) .build() .await; @@ -94,7 +95,7 @@ async fn test_event_table_schema_reconciliation_with_auto_create_disabled() { event_types: Arc::new(EventTypeFilter::Only( [String::from("create_database")].into_iter().collect(), )), - ..Default::default() + ..test_event_recorder_options() }) .build() .await; @@ -185,6 +186,7 @@ async fn test_database_ddl_events() { let home_dir = create_temp_dir("test_database_ddl_events_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_database_ddl_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_frontend_auto_create_table(false) .with_store_config(store_config) @@ -208,6 +210,7 @@ async fn test_database_ddl_events() { async fn test_standalone_database_ddl_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_database_ddl_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let instance = standalone.fe_instance(); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 92888cced3c..8e3057bb7aa 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -228,7 +228,7 @@ fn database_export_request(directory: &std::path::Path) -> table::requests::Copy } } -async fn database_export_roundtrip(instance: &Arc) { +async fn database_export_roundtrip(instance: &Arc, parallelism: usize) { let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); let (first_logical_table_names, _, renamed_physical_table) = create_metric_export_source_tables(instance, "db_a", "dense").await; @@ -253,7 +253,9 @@ async fn database_export_roundtrip(instance: &Arc) { ]; let mut names = selected.clone(); names.extend([renamed_physical_table, "dashboard".into()]); - let req = database_export_request(&destination.path().join("data")); + let mut req = database_export_request(&destination.path().join("data")); + req.with + .insert("parallelism".into(), parallelism.to_string()); let executor = instance.statement_executor(); let captured = executor .capture_database_export_tables(&req, None, &QueryContext::arc()) @@ -448,7 +450,7 @@ async fn database_export_standalone_roundtrip() { let standalone = GreptimeDbStandaloneBuilder::new("database_export") .build() .await; - database_export_roundtrip(standalone.fe_instance()).await; + database_export_roundtrip(standalone.fe_instance(), 1).await; } #[tokio::test(flavor = "multi_thread")] @@ -464,7 +466,7 @@ async fn database_export_distributed_roundtrip() { ) .build(false) .await; - database_export_roundtrip(cluster.fe_instance()).await; + database_export_roundtrip(cluster.fe_instance(), 4).await; } #[tokio::test(flavor = "multi_thread")] @@ -1414,6 +1416,59 @@ async fn export_http(instance: Arc) -> (String, tokio::task::JoinHandl (addr, task) } +#[derive(Default)] +struct ImportRequests { + batches: std::sync::atomic::AtomicUsize, + fail_batch: std::sync::atomic::AtomicBool, + capabilities: std::sync::atomic::AtomicUsize, +} + +async fn import_http( + instance: Arc, + older: bool, +) -> (String, tokio::task::JoinHandle<()>, Arc) { + use axum::response::IntoResponse; + use std::sync::atomic::Ordering; + let requests = Arc::new(ImportRequests::default()); + requests.fail_batch.store(!older, Ordering::SeqCst); + let observed = requests.clone(); + let server = servers::http::HttpServerBuilder::new(Default::default()) + .with_sql_handler(instance) + .build(); + let app = server + .build(server.make_app()) + .unwrap() + .layer(axum::middleware::from_fn( + move |request: axum::extract::Request, next: axum::middleware::Next| { + let observed = observed.clone(); + async move { + let batch = request.uri().path() == "/v1/ddl/logical-tables"; + if request.uri().path() == "/v1/capabilities" { + observed.capabilities.fetch_add(1, Ordering::SeqCst); + if older { + return axum::http::StatusCode::NOT_FOUND.into_response(); + } + } + if batch { + observed.batches.fetch_add(1, Ordering::SeqCst); + } + let response = next.run(request).await; + if batch && observed.fail_batch.swap(false, Ordering::SeqCst) { + assert!(response.status().is_success()); + return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + response + } + }, + )); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (addr, task, requests) +} + struct FailSecondChunk { fail: std::sync::atomic::AtomicBool, copies: std::sync::atomic::AtomicUsize, @@ -1449,16 +1504,69 @@ impl servers::interceptor::SqlQueryInterceptor for FailSecondChunk { #[tokio::test] async fn metric_export_v2_cli_resume_roundtrip() { - metric_export_v2_cli_roundtrip(false).await; + metric_export_v2_cli_roundtrip(false, false, false).await; } #[tokio::test] -#[ignore = "requires GT_S3_* credentials and an S3-compatible test bucket"] -async fn metric_export_v2_cli_s3_resume_roundtrip() { - metric_export_v2_cli_roundtrip(true).await; +async fn metric_import_v2_older_target_roundtrip() { + metric_export_v2_cli_roundtrip(false, false, true).await; } -async fn metric_export_v2_cli_roundtrip(s3: bool) { +#[tokio::test] +async fn metric_export_v2_cli_s3_resume_roundtrip() { + // The CLI takes an explicit `--s3-endpoint`, so this only runs against an + // S3-compatible endpoint (MinIO in CI), not against bucket+region alone. + if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) + && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) + { + metric_export_v2_cli_roundtrip(true, false, false).await; + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn packed_export_v2_cli_local_roundtrip() { + metric_export_v2_cli_roundtrip(false, true, false).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn packed_export_v2_cli_s3_roundtrip() { + if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) + && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) + { + metric_export_v2_cli_roundtrip(true, true, false).await; + } +} + +async fn exported_table_path( + store: &object_store::ObjectStore, + packed: bool, + chunk: u32, + name: &str, +) -> String { + let prefix = format!("data/public/{chunk}/"); + if !packed { + return format!("{prefix}{name}.parquet"); + } + let index: common_datasource::packed_snapshot::PackIndex = serde_json::from_slice( + &store + .read(&format!("{prefix}pack-index.json")) + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + format!( + "{prefix}{}", + index + .tables + .iter() + .find(|t| t.table_name == name) + .unwrap() + .object + ) +} + +async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { use servers::interceptor::SqlQueryInterceptorRef; let plugins = common_base::Plugins::new(); let faults = Arc::new(FailSecondChunk { @@ -1491,7 +1599,24 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { .await; sql(instance, "INSERT INTO audit VALUES ('a',1,1),('z',NULL,3)").await; sql(instance, "CREATE VIEW dashboard AS SELECT * FROM audit").await; + let special_name = "audit.dashboard"; + let quoted_special = special_name.replace('"', "\"\""); + sql( + instance, + &format!("CREATE VIEW \"{quoted_special}\" AS SELECT * FROM audit"), + ) + .await; + if packed && !s3 { + for id in 0..128 { + let name = format!("batch_{id:03}"); + sql(instance, &format!("CREATE TABLE {name} (host STRING PRIMARY KEY, val DOUBLE, ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='v2_a')")).await; + names.push(name); + } + } sql(instance, "CREATE DATABASE z_later").await; + if packed { + sql(instance, "CREATE DATABASE empty_schema").await; + } sql( instance, "CREATE TABLE z_later.events (ts TIMESTAMP TIME INDEX, val BIGINT)", @@ -1520,9 +1645,27 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { names.push("bulk".into()); } let (addr, server) = export_http(instance.clone()).await; + if packed && !s3 { + let client = cli::DatabaseClient::new( + addr.clone(), + "greptime".into(), + None, + std::time::Duration::from_secs(60), + None, + true, + ); + let error = client.sql_in_public( + "SHOW CREATE TABLE audit; SHOW CREATE TABLE missing_middle; SHOW CREATE VIEW dashboard" + ).await.unwrap_err(); + assert!(error.to_string().contains("SQL request failed"), "{error}"); + } + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); for experimental in [false, true] { for layout in ["packed", "invalid"] { + if experimental && layout == "packed" { + continue; + } let path = destination .path() .join(format!("unsupported-{experimental}-{layout}")); @@ -1537,7 +1680,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { ) .await .remove(0); - let error = result.expect_err("export must reject import-only layouts"); + let error = result.expect_err("export must reject unsupported layouts"); assert!( format!("{error:?}").contains("metric_data_layout"), "{error:?}" @@ -1545,6 +1688,41 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { assert!(!path.exists()); } } + if packed && !s3 { + for window in [None, Some("3ms")] { + let path = destination.path().join("empty-range"); + let uri = url::Url::from_directory_path(&path).unwrap(); + let mut args = vec![ + "export-v2", + "create", + "--addr", + &addr, + "--to", + uri.as_str(), + "--schemas", + "public", + "--experimental-metric-export", + "--metric-data-layout", + "packed", + "--no-proxy", + "--start-time", + "1970-01-01T00:00:00Z", + "--end-time", + "1970-01-01T00:00:00Z", + ]; + if let Some(window) = window { + args.extend(["--chunk-time-window", window]); + } + let error = run_data_cli(&args).await.unwrap_err(); + assert!( + error + .to_string() + .contains("Packed export requires --start-time to be earlier than --end-time"), + "{error}" + ); + assert!(!path.exists()); + } + } let (uri, store, storage_args) = if s3 { let endpoint = std::env::var("GT_S3_ENDPOINT_URL").unwrap(); let bucket = std::env::var("GT_S3_BUCKET").unwrap(); @@ -1597,7 +1775,11 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { "--to", &uri, "--schemas", - "public,z_later", + if packed { + "public,z_later,empty_schema" + } else { + "public,z_later" + }, "--experimental-metric-export", "--no-proxy", "--start-time", @@ -1610,6 +1792,9 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { "never", ]; args.extend(storage_args.iter().map(String::as_str)); + if packed { + args.extend(["--metric-data-layout", "packed"]); + } assert!(run_data_cli(&args).await.is_err()); let before: cli::export_v2::manifest::Manifest = serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); @@ -1624,7 +1809,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { if s3 { assert!( store - .stat("data/public/1/bulk.parquet") + .stat(&exported_table_path(&store, packed, 1, "bulk").await) .await .unwrap() .content_length() @@ -1635,7 +1820,12 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { for path in &before.chunks[0].files { preserved.push((path.clone(), store.read(path).await.unwrap().to_vec())); } - assert!(store.exists("data/public/2/audit.parquet").await.unwrap()); + assert!( + store + .exists(&exported_table_path(&store, packed, 2, "audit").await) + .await + .unwrap() + ); store .write("data/public/2/unknown.txt", "keep") .await @@ -1659,6 +1849,59 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { let after: cli::export_v2::manifest::Manifest = serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); assert!(after.is_complete()); + assert_eq!(after.version, if packed { 2 } else { 1 }); + assert_eq!(before.data_layout, after.data_layout); + if packed { + use common_datasource::packed_snapshot::{ObjectKind, PackIndex}; + for chunk in &after.chunks { + let empty: PackIndex = serde_json::from_slice( + &store + .read(&format!("data/empty_schema/{}/pack-index.json", chunk.id)) + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + empty.validate_membership([]).unwrap(); + let prefix = format!("data/public/{}/", chunk.id); + let index: PackIndex = serde_json::from_slice( + &store + .read(&format!("{prefix}pack-index.json")) + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + index + .validate_membership(names.iter().map(String::as_str)) + .unwrap(); + assert_eq!( + index + .objects + .iter() + .filter(|o| o.kind == ObjectKind::Pack) + .count(), + 1 + ); + assert!(index.tables.iter().filter(|t| t.row_count == 0).count() >= 2); + let mut expected = index + .objects + .iter() + .map(|o| format!("{prefix}{}", o.path)) + .collect::>(); + expected.push(format!("{prefix}pack-index.json")); + expected.sort(); + assert_eq!( + chunk + .files + .iter() + .filter(|f| f.starts_with(&prefix)) + .cloned() + .collect::>(), + expected + ); + } + } assert_eq!( serde_json::to_value(&before.chunks[0]).unwrap(), serde_json::to_value(&after.chunks[0]).unwrap() @@ -1679,7 +1922,8 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { ) .await; } - let (target_addr, target_server) = export_http(target.fe_instance().clone()).await; + let (target_addr, target_server, requests) = + import_http(target.fe_instance().clone(), older).await; let state = destination.path().join("restore-state.json"); let mut import = vec![ "import-v2", @@ -1694,7 +1938,29 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { "never", ]; import.extend(storage_args.iter().map(String::as_str)); + use std::sync::atomic::Ordering; + if !older { + assert!(run_data_cli(&import).await.is_err()); + assert_eq!(requests.batches.load(Ordering::SeqCst), 1); + assert_eq!(requests.capabilities.load(Ordering::SeqCst), 1); + let failed: serde_json::Value = + serde_json::from_slice(&std::fs::read(&state).unwrap()).unwrap(); + assert_eq!(failed["ddl_completed"], false); + assert!( + failed["tasks"] + .as_array() + .unwrap() + .iter() + .all(|t| t["status"] == "pending") + ); + } run_data_cli(&import).await.unwrap(); + assert_eq!( + requests.capabilities.load(Ordering::SeqCst), + if older { 1 } else { 2 } + ); + assert_eq!(requests.batches.load(Ordering::SeqCst) > 1, !older); + assert!(!state.exists()); for name in names { let source_table = table(instance, &name).await; let target_table = table(target.fe_instance(), &name).await; @@ -1720,6 +1986,92 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { ) .await ); + for view in ["dashboard", special_name] { + let query = format!( + "SELECT * FROM \"{}\" ORDER BY ts,host", + view.replace('"', "\"\"") + ); + assert_eq!( + values(instance, &query).await, + values(target.fe_instance(), &query).await + ); + assert_eq!( + table(instance, view).await.schema().column_schemas(), + table(target.fe_instance(), view) + .await + .schema() + .column_schemas() + ); + } + if packed { + let schema_uri = format!("{uri}/schema-only"); + let mut schema_args = vec![ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &schema_uri, + "--schemas", + "public", + "--schema-only", + "--experimental-metric-export", + "--metric-data-layout", + "packed", + "--no-proxy", + "--progress", + "never", + ]; + schema_args.extend(storage_args.iter().map(String::as_str)); + run_data_cli(&schema_args).await.unwrap(); + let schema_manifest: cli::export_v2::manifest::Manifest = serde_json::from_slice( + &store + .read("schema-only/manifest.json") + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + assert!(schema_manifest.schema_only && schema_manifest.chunks.is_empty()); + assert_eq!(schema_manifest.version, 1); + assert!(schema_manifest.data_layout.is_none()); + let schema_target = GreptimeDbStandaloneBuilder::new("schema_only_target") + .build() + .await; + let (schema_addr, schema_server, schema_requests) = + import_http(schema_target.fe_instance().clone(), false).await; + schema_requests.fail_batch.store(false, Ordering::SeqCst); + let schema_state = destination.path().join("schema-only-state.json"); + let mut schema_import = vec![ + "import-v2", + "--addr", + &schema_addr, + "--from", + &schema_uri, + "--state-path", + schema_state.to_str().unwrap(), + "--no-proxy", + "--progress", + "never", + ]; + schema_import.extend(storage_args.iter().map(String::as_str)); + let probes = schema_requests.capabilities.load(Ordering::SeqCst); + schema_import.push("--dry-run"); + run_data_cli(&schema_import).await.unwrap(); + assert_eq!(schema_requests.capabilities.load(Ordering::SeqCst), probes); + schema_import.pop(); + let batches = schema_requests.batches.load(Ordering::SeqCst); + run_data_cli(&schema_import).await.unwrap(); + assert!(schema_requests.batches.load(Ordering::SeqCst) > batches); + assert!(!schema_state.exists()); + assert!( + values(schema_target.fe_instance(), "SELECT * FROM dashboard") + .await + .is_empty() + ); + schema_server.abort(); + let _ = schema_server.await; + } server.abort(); target_server.abort(); if s3 { diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index 122c36c84f4..9d2884468d2 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -15,7 +15,9 @@ use std::sync::Arc; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::setup_authenticated_grpc_database; +use tests_integration::test_util::{ + setup_authenticated_grpc_database, test_event_recorder_options, +}; use uuid::Uuid; use crate::event_recorder_test_util::{ @@ -31,6 +33,7 @@ const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; async fn test_standalone_flow_ddl_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_flow_ddl_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let instance = standalone.fe_instance(); diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 4721a442a80..48749b784a0 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -43,6 +43,9 @@ async fn test_batch_gc_event() { common_telemetry::init_default_ut_logging(); let (store_config, _guard) = get_test_store_config(&store_type); let home_dir = create_temp_dir("test_batch_gc_event_data_home"); + // Keeps the production event-recorder flush interval: `assert_sst_count` + // counts SST files across every table, so eagerly flushed event rows would + // show up in the count. let cluster = GreptimeDbClusterBuilder::new("test_batch_gc_event") .await .with_datanodes(1) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 025d56b959a..e9cfee7e9d4 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Retain coverage of the legacy constructors, including TLS configuration. +#![allow(deprecated)] + use std::sync::Arc; use api::v1::alter_table_expr::Kind; @@ -19,12 +22,13 @@ 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::query_request::Query; use api::v1::value::ValueData; use api::v1::{ AddColumn, AddColumns, AlterTableExpr, Basic, Column, ColumnDataType, ColumnDef, ColumnSchema, CreateTableExpr, GreptimeRequest, InsertRequest, InsertRequests, PromInstantQuery, - PromRangeQuery, PromqlRequest, RequestHeader, Row, RowInsertRequest, RowInsertRequests, Rows, - SemanticType, Value, column, + PromRangeQuery, PromqlRequest, QueryRequest, RequestHeader, Row, RowInsertRequest, + RowInsertRequests, Rows, SemanticType, Value, column, }; use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; @@ -72,6 +76,7 @@ use tests_integration::test_util::{ setup_grpc_server_with_auto_create_table_disabled, setup_grpc_server_with_user_provider, }; use tonic::Request; +use tonic::codec::CompressionEncoding; use tonic::metadata::MetadataValue; use crate::both_deployment_cases; @@ -89,9 +94,13 @@ macro_rules! grpc_test { async fn [< $test >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { - let _ = $crate::grpc::$test(store_type).await; + // Support both unit tests and fallible tests without discarding errors. + let result = $crate::grpc::$test(store_type).await; + assert_eq!( + std::process::Termination::report(result), + std::process::ExitCode::SUCCESS, + ); } - } )* } @@ -127,6 +136,7 @@ macro_rules! grpc_tests { test_grpc_timezone, test_grpc_tls_config, test_grpc_memory_limit, + test_grpc_compressed_memory_reservation, ); )* }; @@ -198,29 +208,128 @@ pub async fn test_grpc_message_size_ok(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +/// Both the server (`servers::grpc::builder`) and the standard client +/// (`configure_tonic_client!`) enable zstd, so a plain round trip never shows +/// whether negotiation actually happened. This drives raw tonic clients instead +/// and asserts on the `grpc-encoding` the server answers with. pub async fn test_grpc_zstd_compression(store_type: StorageType) { - // server and client both support gzip - let config = GrpcServerConfig { - max_recv_message_size: 1024, - max_send_message_size: 1024, - ..Default::default() - }; - let (_db, fe_grpc_server) = setup_grpc_server_with( - store_type, - "test_grpc_zstd_compression", - None, - Some(config), - None, - ) - .await; + let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_grpc_zstd_compression").await; let addr = fe_grpc_server.bind_addr().unwrap().to_string(); - let grpc_client = Client::with_urls(vec![addr]); + let ddl = |sql: &str| GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::Query(QueryRequest { + query: Some(Query::Sql(sql.to_string())), + })), + }; + + // Sends zstd and accepts zstd: the server has to decode a compressed request + // body and compress its response. + let mut zstd_client = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd) + .accept_compressed(CompressionEncoding::Zstd); + let response = zstd_client + .handle(Request::new(ddl( + "CREATE TABLE zstd_compression (ts TIMESTAMP TIME INDEX, payload STRING)", + ))) + .await + .unwrap(); + assert_eq!( + response + .metadata() + .get("grpc-encoding") + .map(|v| v.to_str().unwrap()), + Some("zstd") + ); + + // A payload far above the zstd frame overhead, so the request body really is + // compressed rather than passed through. + let payload = "compressible-".repeat(4096); + let insert = |ts: i64| GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "zstd_compression".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() + }, + ColumnSchema { + column_name: "payload".to_string(), + semantic_type: SemanticType::Field as i32, + datatype: ColumnDataType::String as i32, + ..Default::default() + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(ts)), + }, + Value { + value_data: Some(ValueData::StringValue(payload.clone())), + }, + ], + }], + }), + }], + })), + }; + + let response = zstd_client + .handle(Request::new(insert(1000))) + .await + .unwrap(); + let ResponseBody::AffectedRows(rows) = response.into_inner().response.unwrap(); + assert_eq!(rows.value, 1); + + // A client that does not advertise zstd gets an uncompressed response. + let mut plain_client = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap(); + let response = plain_client + .handle(Request::new(insert(2000))) + .await + .unwrap(); + assert!(response.metadata().get("grpc-encoding").is_none()); + + // Both payloads survived their respective paths intact. let db = Database::new_with_dbname( format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), - grpc_client, + Client::with_urls(vec![addr]), ); - db.sql("show tables;").await.unwrap(); + let sql = format!( + "SELECT count(*) AS c FROM zstd_compression WHERE length(payload) = {}", + payload.len() + ); + let OutputData::Stream(stream) = db.sql(&sql).await.unwrap().data else { + panic!("expected a stream"); + }; + let recordbatches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + recordbatches.pretty_print().unwrap(), + "\ ++---+ +| c | ++---+ +| 2 | ++---+" + ); + let _ = fe_grpc_server.shutdown().await; } @@ -250,6 +359,104 @@ pub async fn test_grpc_message_size_limit_send(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +/// Transport-compressed gRPC requests reserve the worst-case decoded message +/// size (`max_recv_message_size`) against the aggregate quota *before* tonic +/// decompresses them. With a quota smaller than that bound, compressed +/// requests must fail fast with `RESOURCE_EXHAUSTED` (before decoding), while +/// uncompressed requests are still admitted through the exact post-decode +/// charge. With the default unlimited limiter, compression keeps working. +pub async fn test_grpc_compressed_memory_reservation(store_type: StorageType) { + use api::v1::query_request::Query; + use api::v1::{GreptimeRequest, QueryRequest}; + use tonic::codec::CompressionEncoding; + + let config = GrpcServerConfig { + max_recv_message_size: 1024 * 1024, + ..Default::default() + }; + // Quota smaller than the decoding bound: compressed requests cannot be + // pre-admitted. + let memory_limiter = ServerMemoryLimiter::new(512 * 1024, OnExhaustedPolicy::Fail); + let (_db, fe_grpc_server) = setup_grpc_server_with( + store_type, + "test_grpc_compressed_memory_reservation", + None, + Some(config), + Some(memory_limiter), + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let db = Database::new_with_dbname( + format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), + Client::with_urls(vec![addr.clone()]), + ); + db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)") + .await + .unwrap(); + + let insert_request = || { + tonic::Request::new(GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::Query(QueryRequest { + query: Some(Query::Sql( + "INSERT INTO grpc_quota VALUES (1000, 1.0)".to_string(), + )), + })), + }) + }; + + // Compressed request: the pre-decode reservation (1 MiB) exceeds the + // 512 KiB quota and must be rejected before any decoding happens. + let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd); + let status = compressed.handle(insert_request()).await.unwrap_err(); + assert_eq!( + status.code(), + tonic::Code::ResourceExhausted, + "compressed request must be pre-admitted against the quota, got: {status}" + ); + + // Uncompressed request on the same server: still admitted (its exact + // serialized size fits the quota). + let mut plain = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap(); + plain.handle(insert_request()).await.unwrap(); + + let _ = fe_grpc_server.shutdown().await; + + // Default (unlimited) limiter: compressed requests keep working. + let (_db, fe_grpc_server) = setup_grpc_server_with( + store_type, + "test_grpc_compressed_memory_reservation_unlimited", + None, + Some(GrpcServerConfig::default()), + None, + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + let db = Database::new_with_dbname( + format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), + Client::with_urls(vec![addr.clone()]), + ); + db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)") + .await + .unwrap(); + let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd); + compressed.handle(insert_request()).await.unwrap(); + let _ = fe_grpc_server.shutdown().await; +} + pub async fn test_grpc_message_size_limit_recv(store_type: StorageType) { let config = GrpcServerConfig { max_recv_message_size: 10, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index e4e30baeba8..91979ddbf59 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -15,18 +15,20 @@ use std::collections::BTreeMap; use std::io::Write; use std::str::FromStr; +use std::sync::Arc; use std::time::Duration; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::metadata::MetricType as RemoteWriteV2MetricType; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Sample as RemoteWriteV2Sample, - TimeSeries as RemoteWriteV2TimeSeries, + BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Request as RemoteWriteV2Request, + Sample as RemoteWriteV2Sample, TimeSeries as RemoteWriteV2TimeSeries, }; use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{ Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, TimeSeries, WriteRequest, }; +use api::v1::RowInsertRequests; use auth::{UserProviderRef, user_provider_from_option}; use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode}; use base64::prelude::{BASE64_STANDARD, Engine as _}; @@ -69,7 +71,9 @@ use servers::http::header::constants::{ }; use servers::http::header::{GREPTIME_DB_HEADER_NAME, GREPTIME_TIMEZONE_HEADER_NAME}; use servers::http::otlp::GoogleRpcStatus; -use servers::http::prometheus::{Column, PrometheusJsonResponse, PrometheusResponse}; +use servers::http::prometheus::{ + Column, PromQueryResult, PrometheusJsonResponse, PrometheusResponse, +}; use servers::http::result::error_result::ErrorResponse; use servers::http::result::greptime_result_v1::GreptimedbV1Response; use servers::http::result::influxdb_result_v1::{InfluxdbOutput, InfluxdbV1Response}; @@ -79,13 +83,17 @@ 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 session::context::QueryContextRef; use standalone::options::StandaloneOptions; use table::table_name::TableName; use tests_integration::test_util::{ - 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, + MockInstanceImpl, StorageType, TestGuard, assert_wal_delta, build_test_prom_server, + setup_test_http_app, setup_test_http_app_with_frontend, + setup_test_http_app_with_frontend_and_custom_options, + 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, + setup_test_prom_app_with_frontend_and_memory_limiter, + setup_test_prom_app_with_frontend_batched, }; use urlencoding::encode; use yaml_rust::YamlLoader; @@ -93,6 +101,31 @@ use yaml_rust::YamlLoader; use crate::both_deployment_cases; use crate::event_recorder_test_util::assert_procedure_actor_by_table; +/// Registers tests that exercise nothing below the HTTP layer, so running them +/// against a remote object store only pays bucket setup and teardown. +#[macro_export] +macro_rules! http_local_tests { + ($($(#[$meta:meta])* $test:ident),*,) => { + mod integration_http_local_test { + $( + #[tokio::test(flavor = "multi_thread")] + $( + #[$meta] + )* + async fn $test() { + let store_type = tests_integration::test_util::StorageType::File; + // Support both unit tests and fallible tests without discarding errors. + let result = $crate::http::$test(store_type).await; + assert_eq!( + std::process::Termination::report(result), + std::process::ExitCode::SUCCESS, + ); + } + )* + } + }; +} + #[macro_export] macro_rules! http_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -132,26 +165,29 @@ macro_rules! http_tests { test_http_analyze_stream_tql, test_http_sql_slow_query, test_prometheus_promql_api, + test_promql_over_non_millisecond_physical_tables, test_prometheus_label_replace_response, test_prom_http_api, - test_metrics_api, - test_health_api, - test_status_api, + test_prom_remote_compression_bomb_rejected, + test_prom_remote_write_decoded_body_charged, + test_influxdb_write_decoded_body_charged, test_config_api, - test_dynamic_tracer_toggle, - test_dashboard_path, test_dashboard_api, test_prometheus_remote_write, test_prometheus_remote_write_v2, test_prometheus_remote_write_v2_native_histogram, test_prometheus_remote_write_batched, + test_prometheus_remote_write_batched_mixed_time_index_units, + test_prometheus_remote_write_batched_interceptor_time_index_units, + test_prometheus_remote_write_v2_batched_interceptor_time_index_units, + test_prometheus_remote_write_batched_microsecond_physical_table, + test_prometheus_remote_write_batched_conflicting_physical_selections, test_prometheus_remote_special_labels, test_prometheus_remote_schema_labels, test_prometheus_remote_write_with_pipeline, test_vm_proto_remote_write, test_pipeline_api, - test_test_pipeline_api, test_pipeline_name_in_header, test_plain_text_ingestion, test_pipeline_auto_transform, @@ -187,8 +223,6 @@ macro_rules! http_tests { test_loki_json_logs_with_pipeline, test_elasticsearch_logs, test_elasticsearch_logs_with_index, - test_splunk_health, - test_splunk_health_is_public, test_splunk_logs, test_splunk_raw, test_log_query, @@ -483,6 +517,459 @@ async fn test_sql_skip_wal(distributed: bool) { check_http_skip_wal("sql", &cases, distributed).await; } +fn logical_ddl(name: &str) -> String { + format!( + "CREATE TABLE IF NOT EXISTS {name} (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (on_physical_table='phy.dot')" + ) +} + +#[tokio::test(flavor = "multi_thread")] +async fn logical_ddl_http_standalone() { + check_logical_ddl_http(false).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn logical_ddl_http_distributed() { + check_logical_ddl_http(true).await; +} + +#[tokio::test(flavor = "multi_thread")] +#[ignore = "manual logical DDL throughput measurement"] +async fn logical_ddl_http_benchmark() { + common_telemetry::init_default_ut_logging(); + let count: usize = std::env::var("GT_LOGICAL_DDL_BENCH_TABLES") + .unwrap_or_else(|_| "1000".to_string()) + .parse() + .unwrap(); + for (name, endpoint, batch_size) in [ + ("single", "/v1/sql", 1), + ("multi_sql", "/v1/sql", 128), + ("batch", "/v1/ddl/logical-tables", 128), + ] { + let mut instance = + MockInstanceImpl::Standalone( + tests_integration::standalone::GreptimeDbStandaloneBuilder::new(&format!( + "logical_ddl_bench_{name}" + )) + .with_event_recorder_options( + tests_integration::test_util::test_event_recorder_options(), + ) + .build() + .await, + ); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(instance.frontend()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let response = execute_sql(&client, "CREATE TABLE \"phy.dot\" (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + let statements = (0..count) + .map(|i| logical_ddl(&format!("bench_{i}"))) + .collect::>(); + let start = std::time::Instant::now(); + for chunk in statements.chunks(batch_size) { + let response = client + .post(endpoint) + .form(&[("sql", chunk.join(";"))]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + assert_eq!( + response.json::().await["output"] + .as_array() + .unwrap() + .len(), + chunk.len() + ); + } + let elapsed = start.elapsed(); + let requests = count.div_ceil(batch_size); + let procedures = if name == "batch" { requests } else { count }; + wait_for_event_data(&client, + "SELECT count(DISTINCT procedure_id), count(DISTINCT table_name) FROM greptime_private.events WHERE type='create_logical_tables' AND json_get_string(procedure_trigger, 'type')='Submitted' AND table_name LIKE 'bench_%'", + &format!("[[{procedures},{count}]]")).await; + common_telemetry::info!( + "Logical DDL: mode={name}, tables={count}, concurrency=1, HTTP requests={requests}, procedures={procedures}, elapsed={elapsed:?}" + ); + instance.shutdown().await; + } +} + +async fn check_logical_ddl_http(distributed: bool) { + common_telemetry::init_default_ut_logging(); + let events = tests_integration::test_util::test_event_recorder_options(); + let mut instance = if distributed { + MockInstanceImpl::Distributed( + tests_integration::cluster::GreptimeDbClusterBuilder::new("batch_logical_ddl") + .await + .with_event_recorder_options(events) + .with_datanodes(3) + .build(true) + .await, + ) + } else { + MockInstanceImpl::Standalone( + tests_integration::standalone::GreptimeDbStandaloneBuilder::new("batch_logical_ddl") + .with_event_recorder_options(events) + .build() + .await, + ) + }; + let frontend = instance.frontend(); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(frontend.clone()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + assert_eq!( + client + .get("/v1/capabilities") + .send() + .await + .json::() + .await["metric_batch_ddl"], + 1 + ); + let response = execute_sql(&client, "CREATE TABLE \"phy.dot\" (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + + for sql in [ + String::new(), + format!("{}; SELECT 1", logical_ddl("rejected")), + format!("{}; CREATE VIEW v AS SELECT 1", logical_ddl("rejected")), + format!( + "{}; CREATE TABLE copy LIKE batch_0", + logical_ddl("rejected") + ), + format!("{}; CREATE TABLE copy AS SELECT 1", logical_ddl("rejected")), + logical_ddl("rejected").replace("WITH (", "WITH (physical_metric_table='', "), + format!( + "{}; CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX)", + logical_ddl("rejected") + ), + format!( + "{};{}", + logical_ddl("rejected"), + logical_ddl("public.rejected") + ), + format!("{};{}", logical_ddl("rejected"), logical_ddl("other.t")), + format!( + "{};{}", + logical_ddl("rejected"), + logical_ddl("other_physical").replace("phy.dot", "other") + ), + (0..129) + .map(|i| logical_ddl(&format!("rejected_{i}"))) + .collect::>() + .join(";"), + ] { + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(!response.status().is_success(), "{}", response.text().await); + } + validate_data( + "batch_rejected", + &client, + "SELECT count(*) FROM information_schema.tables WHERE table_name LIKE 'rejected%'", + "[[0]]", + ) + .await; + + let sql = (0..128) + .map(|i| logical_ddl(&format!("batch_{i}"))) + .collect::>() + .join(";"); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + assert_eq!( + response.json::().await["output"], + json!(vec![json!({"affectedrows": 0}); 128]) + ); + wait_for_event_data(&client, + "SELECT count(DISTINCT procedure_id), count(DISTINCT table_name) FROM greptime_private.events WHERE type='create_logical_tables' AND json_get_string(procedure_trigger, 'type')='Submitted' AND table_name LIKE 'batch_%'", + "[[1,128]]").await; + + let physical_id = frontend + .catalog_manager() + .table("greptime", "public", "phy.dot", None) + .await + .unwrap() + .unwrap() + .table_info() + .table_id(); + let mut ids = std::collections::HashSet::new(); + for i in 0..128 { + let table = frontend + .catalog_manager() + .table("greptime", "public", &format!("batch_{i}"), None) + .await + .unwrap() + .unwrap(); + let id = table.table_info().table_id(); + assert!(ids.insert(id)); + assert_ne!(id, physical_id); + let bound_physical_id = frontend + .table_metadata_manager() + .table_route_manager() + .get_physical_table_id(id) + .await + .unwrap(); + assert_eq!(bound_physical_id, physical_id); + assert_eq!( + table + .table_info() + .meta + .options + .extra_options + .get("on_physical_table") + .map(String::as_str), + Some("phy.dot") + ); + assert_eq!(table.schema().column_schemas().len(), 3); + } + let sql = logical_ddl("one_at_byte_limit"); + let padded = format!("{sql}{}", " ".repeat(1024 * 1024 - sql.len())); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", padded)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + + let response = execute_sql(&client, "CREATE DATABASE inherited WITH ('skip_wal'='true', ttl='1h', 'compaction.type'='twcs'); CREATE TABLE inherited.phy (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + let inherited = logical_ddl("inherited.ordinary").replace("phy.dot", "phy"); + let response = execute_sql(&client, &inherited).await; + assert!(response.status().is_success(), "{}", response.text().await); + let sql = format!( + "{};{}", + logical_ddl("batched").replace("phy.dot", "phy"), + logical_ddl("explicit") + .replace("phy.dot", "phy") + .replace("WITH (", "WITH (skip_wal='false', ") + ); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql.as_str()), ("db", "greptime-inherited")]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + let catalog = frontend.catalog_manager(); + let ordinary = catalog + .table("greptime", "inherited", "ordinary", None) + .await + .unwrap() + .unwrap(); + let batched = catalog + .table("greptime", "inherited", "batched", None) + .await + .unwrap() + .unwrap(); + let explicit = catalog + .table("greptime", "inherited", "explicit", None) + .await + .unwrap() + .unwrap(); + assert_eq!( + ordinary.table_info().meta.options, + batched.table_info().meta.options + ); + assert!(batched.table_info().meta.options.skip_wal); + assert!(!explicit.table_info().meta.options.skip_wal); + assert!( + !batched + .table_info() + .meta + .options + .extra_options + .contains_key("compaction.type") + ); + + let response = execute_sql(&client, "CREATE TABLE partitioned (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) PARTITION ON COLUMNS (host) (host < 'm', host >= 'm') ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + let partitioned = logical_ddl("bad_partition") + .replace("phy.dot", "partitioned") + .replace( + "ENGINE=", + "PARTITION ON COLUMNS (host) (host < 'n', host >= 'n') ENGINE=", + ); + let response = execute_sql(&client, &partitioned).await; + assert!(!response.status().is_success()); + let sql = format!( + "{};{partitioned}", + logical_ddl("before_bad_partition").replace("phy.dot", "partitioned") + ); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(!response.status().is_success()); + assert!( + catalog + .table("greptime", "public", "before_bad_partition", None) + .await + .unwrap() + .is_none() + ); + + let matching = partitioned + .replace("bad_partition", "matching_partition") + .replace("'n'", "'m'"); + let response = execute_sql(&client, &matching).await; + assert!(response.status().is_success(), "{}", response.text().await); + let sql = format!( + "{};{}", + matching.replace("matching_partition", "batch_matching_partition"), + logical_ddl("batch_inherited_partition").replace("phy.dot", "partitioned") + ); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + + let original_id = frontend + .catalog_manager() + .table("greptime", "public", "batch_0", None) + .await + .unwrap() + .unwrap() + .table_info() + .table_id(); + let sql = format!("{};{}", logical_ddl("batch_0"), logical_ddl("batch_new")); + for _ in 0..2 { + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", &sql)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + } + let old = frontend + .catalog_manager() + .table("greptime", "public", "batch_0", None) + .await + .unwrap() + .unwrap(); + assert_eq!(original_id, old.table_info().table_id()); + let response = execute_sql( + &client, + "INSERT INTO batch_0 (ts, host, val) VALUES (1, 'a', 2.0); INSERT INTO batch_new (ts, host, val) VALUES (2, 'b', 3.0)", + ) + .await; + assert!(response.status().is_success(), "{}", response.text().await); + validate_data( + "batch_values", + &client, + "SELECT val FROM batch_0 UNION ALL SELECT val FROM batch_new ORDER BY val", + "[[2.0],[3.0]]", + ) + .await; + let ctx = session::context::QueryContext::arc(); + let provider = + auth::static_user_provider_from_option("static_user_provider:cmd:readonly:ro=password") + .unwrap(); + let user = auth::UserProvider::authenticate( + &provider, + auth::Identity::UserId("readonly", None), + auth::Password::PlainText("password".to_string().into()), + ) + .await + .unwrap(); + ctx.set_current_user(user); + frontend + .plugins() + .get_or_insert::(auth::DefaultPermissionChecker::arc); + let result = servers::query_handler::sql::SqlQueryHandler::create_logical_tables( + frontend.as_ref(), + &logical_ddl("read_only_denied"), + ctx, + ) + .await; + assert!(result.is_err()); + assert!( + frontend + .catalog_manager() + .table("greptime", "public", "read_only_denied", None) + .await + .unwrap() + .is_none() + ); + + struct FailAfterExecution(&'static str); + impl servers::interceptor::SqlQueryInterceptor for FailAfterExecution { + type Error = frontend::error::Error; + fn post_execute( + &self, + _: common_query::Output, + _: QueryContextRef, + ) -> frontend::error::Result { + use common_query::Output; + match self.0 { + "rows" => Ok(Output::new_with_affected_rows(1)), + "records" => Ok(Output::new_with_record_batches( + common_recordbatch::RecordBatches::empty(), + )), + "stream" => Ok(Output::new_with_stream(Box::pin( + common_recordbatch::RecordBatchStreamWrapper::new( + Arc::new(datatypes::schema::Schema::new(vec![])), + futures::stream::iter([Err::( + common_recordbatch::error::CreateRecordBatchesSnafu { + reason: "sensitive-option-value", + } + .build(), + )]), + ), + ))), + _ => frontend::error::InvalidSqlSnafu { + err_msg: "sensitive-option-value", + } + .fail(), + } + } + } + frontend + .plugins() + .insert::>(Arc::new( + FailAfterExecution("error"), + )); + for mode in ["error", "rows", "records", "stream"] { + frontend + .plugins() + .map_mut::, _, _>( + |plugin| *plugin.unwrap() = Arc::new(FailAfterExecution(mode)), + ); + let name = format!("created_before_error_{mode}"); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", logical_ddl(&name))]) + .send() + .await; + assert!(!response.status().is_success(), "{mode}"); + assert!(!format!("{:?}", response.headers()).contains("sensitive-option-value")); + let body = response.json::().await; + assert!(body.get("output").is_none()); + assert_eq!(body["error"], "logical-table batch failed"); + assert!( + frontend + .catalog_manager() + .table("greptime", "public", &name, None) + .await + .unwrap() + .is_some() + ); + } + instance.shutdown().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; @@ -988,6 +1475,31 @@ pub async fn test_prometheus_label_replace_response(store_type: StorageType) { .unwrap() ); + // A series whose source value does not match keeps its labels, without `host_copy`. + let query = encode(r#"label_replace(demo, "host_copy", "$1", "host", "other(.*)")"#); + let res = client + .get(&format!("/v1/prometheus/api/v1/query?query={query}&time=0")) + .send() + .await; + + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!(body.status, "success"); + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "vector", + "result": [{ + "metric": { + "__name__": "demo", + "host": "host1" + }, + "value": [0.0, "1.1"] + }] + })) + .unwrap() + ); + guard.remove_all().await; } @@ -1696,6 +2208,395 @@ pub async fn test_prom_http_api(store_type: StorageType) { guard.remove_all().await; } +/// A tiny compressed body must not be able to demand an unbounded decoded +/// allocation on the Prometheus remote-storage endpoints. The five-byte raw +/// Snappy varint (`ff ff ff ff 0f`) declares a ~4 GiB decoded length and must +/// be rejected by the decoded-size cap. +pub async fn test_prom_remote_compression_bomb_rejected(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_prom_app_with_frontend(store_type, "prom_remote_bomb").await; + let client = TestClient::new(app).await; + + // Remote read with the five-byte varint bomb. + let res = client + .post("/v1/prometheus/read") + .body(Vec::from([0xffu8, 0xff, 0xff, 0xff, 0x0f])) + .send() + .await; + assert_eq!(res.status(), StatusCode::BAD_REQUEST); + + // Remote write v1 with the same bomb, via both codecs (exercising the + // fallback: the zstd attempt fails on malformed input and the snappy + // fallback must reject the declared size). + for encoding in ["snappy", "zstd"] { + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", encoding) + .body(Vec::from([0xffu8, 0xff, 0xff, 0xff, 0x0f])) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::BAD_REQUEST, + "bomb with Content-Encoding {encoding} must be rejected" + ); + } + + // The server must still serve valid requests afterwards. + let write_request = WriteRequest { + timeseries: prom_store::mock_timeseries(), + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + guard.remove_all().await; +} + +/// The decoded (not encoded) size of a Prometheus remote write body must be +/// charged against the aggregate request-memory quota. +pub async fn test_prom_remote_write_decoded_body_charged(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + // 4 KiB aggregate quota with fail-fast policy. + let memory_limiter = ServerMemoryLimiter::new(4096, OnExhaustedPolicy::Fail); + let (app, mut guard) = setup_test_prom_app_with_frontend_and_memory_limiter( + store_type, + "prom_remote_write_quota", + Some(memory_limiter), + ) + .await; + let client = TestClient::new(app).await; + + // A valid write request whose decoded protobuf clearly exceeds the quota + // while its Snappy wire body stays far below it. + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: "__name__".to_string(), + value: "quota_metric".to_string(), + }, + Label { + name: "payload".to_string(), + // ~8 KiB decoded; compresses to a few hundred bytes. + value: "p".repeat(8 * 1024), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + assert!( + compressed.len() < 2048, + "wire body must stay under the quota" + ); + + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + let status = res.status(); + let body = res.text().await; + assert_eq!( + status, + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the aggregate quota must be rejected, got body: {body}" + ); + + guard.remove_all().await; +} + +/// Decompressed ingestion bodies (routes using `RequestDecompressionLayer`) +/// must be charged against the aggregate request-memory quota as they are +/// decoded, not only their compressed wire size. +pub async fn test_influxdb_write_decoded_body_charged(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + // 4 KiB aggregate quota with fail-fast policy. + let memory_limiter = ServerMemoryLimiter::new(4096, OnExhaustedPolicy::Fail); + let http_opts = servers::http::HttpOptions { + addr: format!("127.0.0.1:{}", common_test_util::ports::get_port()), + ..Default::default() + }; + let (app, mut guard) = setup_test_http_app_with_frontend_and_custom_options( + store_type, + "test_influxdb_decoded_quota", + None, + Some(http_opts), + Some(memory_limiter), + ) + .await; + let client = TestClient::new(app).await; + + // A gzip member expanding to ~8 KiB of line protocol from a tiny wire body. + let lines: Vec = (0..200) + .map(|i| format!("decoded_quota,host=h{} v=1.0,ts={}", i % 20, 1700000000 + i)) + .collect(); + let decoded = lines.join("\n"); + assert!(decoded.len() > 4096); + let compressed = zstd::stream::encode_all(decoded.as_bytes(), 3).unwrap(); + assert!( + compressed.len() < 2048, + "wire body must stay under the quota" + ); + + // Compressed write whose decoded body exceeds the quota: rejected while + // streaming (the extractor surfaces the aborted body as 400). + let res = client + .post("/v1/influxdb/write?db=public") + .header("Content-Type", "text/plain") + .header("Content-Encoding", "zstd") + .body(compressed) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the aggregate quota must be rejected with 429" + ); + + // The server still serves a small compressed write within the quota. + let small = "decoded_quota,host=ok v=1.0,ts=1700000000"; + let compressed_small = zstd::stream::encode_all(small.as_bytes(), 3).unwrap(); + let res = client + .post("/v1/influxdb/write?db=public") + .header("Content-Type", "text/plain") + .header("Content-Encoding", "zstd") + .body(compressed_small) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::NO_CONTENT, + "small compressed write must succeed, got: {}", + res.text().await + ); + + guard.remove_all().await; +} + +/// PromQL must behave identically regardless of the physical metric table's +/// time index unit: the same samples are remote-written into physical tables +/// pre-created with second/micro/nano time indexes (plus a millisecond +/// baseline), and every PromQL response must match the baseline exactly. +/// Regression test for . +pub async fn test_promql_over_non_millisecond_physical_tables(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_prom_app_with_frontend(store_type, "promql_non_ms_units").await; + let client = TestClient::new(app).await; + + let units = [ + ("ms", "timestamp(3)"), + ("second", "timestamp(0)"), + ("micro", "timestamp(6)"), + ("nano", "timestamp(9)"), + ]; + + for (suffix, ts_type) in units { + let db = format!("promql_units_{suffix}"); + let res = client + .get(&format!("/v1/sql?db=public&sql=create database {db}")) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "create database {db}"); + + let res = client + .get(&format!( + "/v1/sql?db={db}&sql=CREATE TABLE greptime_physical_table \ + (greptime_timestamp {ts_type} NOT NULL, greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + )) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::OK, + "create physical table with {ts_type} in {db}" + ); + + // unit_gauge{job="demo"} = 1.0 @ 1000ms, 2.0 @ 2000ms + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "unit_gauge".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![ + Sample { + value: 1.0, + timestamp: 1000, + }, + Sample { + value: 2.0, + timestamp: 2000, + }, + ], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let write_url = format!("/v1/prometheus/write?db={db}"); + let res = client + .post(write_url.as_str()) + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::NO_CONTENT, + "remote write into {db}" + ); + } + + let query_paths = [ + "query?query=unit_gauge&time=2", + "query_range?query=unit_gauge&start=0&end=5&step=1", + "query?query=avg_over_time(unit_gauge[1m])&time=2", + "query?query=rate(unit_gauge[1m])&time=2", + "query?query=unit_gauge[1m]&time=2", + ]; + for path in query_paths { + let mut baseline: Option<(String, PrometheusResponse)> = None; + for (suffix, _) in units { + let db = format!("promql_units_{suffix}"); + let res = client + .get(&format!("/v1/prometheus/api/v1/{path}&db={db}")) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "promql `{path}` on {db}"); + let body = res.json::().await; + assert_eq!(body.status, "success", "promql `{path}` on {db}"); + match &baseline { + None => baseline = Some((suffix.to_string(), body.data)), + Some((baseline_suffix, baseline_data)) => { + assert_eq!( + &body.data, baseline_data, + "promql `{path}` on {db} differs from the {baseline_suffix} baseline" + ); + } + } + } + } + + // Pin the baseline responses so the cross-unit equality above cannot pass + // vacuously on wrong data: the samples are 1.0@1s and 2.0@2s, and + // `rate(unit_gauge[1m])@2s` follows Prometheus's extrapolated rate: + // Δv=1 over a 1s sampled interval extrapolated to 1.5s, divided by the + // 60s window -> 1.5/60 = 0.025 per second. + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=unit_gauge&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "vector", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "value": [2.0, "2.0"] + }] + })) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/query_range?db=promql_units_ms&query=unit_gauge&start=0&end=5&step=1") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "matrix", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "values": [ + [1.0, "1.0"], [2.0, "2.0"], [3.0, "2.0"], [4.0, "2.0"], [5.0, "2.0"] + ] + }] + })) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=unit_gauge[1m]&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "matrix", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "values": [[1.0, "1.0"], [2.0, "2.0"]] + }] + })) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=avg_over_time(unit_gauge[1m])&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "vector", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "value": [2.0, "1.5"] + }] + })) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=rate(unit_gauge[1m])&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + let rate = match body.data { + PrometheusResponse::PromData(data) => match data.result { + PromQueryResult::Vector(v) => v[0].value.as_ref().unwrap().1.clone(), + other => panic!("expected vector, got {other:?}"), + }, + other => panic!("expected prom data, got {other:?}"), + }; + assert_eq!(rate, "0.025"); + + guard.remove_all().await; +} + pub async fn test_metrics_api(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = setup_test_http_app(store_type, "metrics_api").await; @@ -1810,15 +2711,21 @@ pub async fn test_splunk_logs(store_type: StorageType) { .await; let client = TestClient::new(app).await; - // Authenticated SQL query (the user-provider harness requires auth on /v1/sql). + // Authenticated SQL statement (the user-provider harness requires auth on + // /v1/sql). Posted as a form so quotes and parentheses survive. async fn query(client: &TestClient, sql: &str) -> String { let res = client - .get(format!("/v1/sql?sql={sql}").as_str()) + .post("/v1/sql") .header("Authorization", basic_auth("greptime_user", "greptime_pwd")) + .form(&[("sql", sql)]) .send() .await; - assert_eq!(res.status(), StatusCode::OK, "query failed: {sql}"); - res.text().await + let status = res.status(); + let body = res.text().await; + assert_eq!(status, StatusCode::OK, "query failed: {sql}: {body}"); + // A failing statement still comes back as 200 with an `error` envelope. + assert!(!body.contains("\"error\""), "query failed: {sql}: {body}"); + body } // HEC `Authorization: Splunk ` + JSON content type. @@ -1969,6 +2876,34 @@ transform: "custom pipeline should have dropped sourcetype (identity would keep it): {create}" ); + // 6b. Ingest into a pre-created table with a fulltext index, then flush: the + // index file is only written on flush, so this is the one path that reads it + // back. `matches_term` is the predicate the index applier recognises. + query( + &client, + "create table splunk_indexed (greptime_timestamp timestamp(9) time index, \ + \"event\" string fulltext index with(backend='bloom'))", + ) + .await; + let res = send_req( + &client, + splunk_headers(), + "/v1/splunk/services/collector/event?table=splunk_indexed", + br#"{"event":"disk pressure on web-01","time":1700000020}"#.to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + query(&client, "ADMIN FLUSH_TABLE('splunk_indexed')").await; + let rows = get_rows_from_output( + &query( + &client, + "select \"event\" from splunk_indexed where matches_term(\"event\", 'pressure')", + ) + .await, + ); + assert_eq!(rows, r#"[["disk pressure on web-01"]]"#); + // 7. Auth failures return HEC codes: missing token -> 2 (401), bad token -> 4 (403). let res = send_req( &client, @@ -2343,19 +3278,6 @@ providers = []"#, ) }; - let vector_index_config = if cfg!(feature = "vector_index") { - r#" -[region_engine.mito.vector_index] -create_on_flush = "auto" -create_on_compaction = "auto" -apply_on_query = "auto" -mem_threshold_on_create = "auto" - -"# - } else { - "\n" - }; - let expected_toml_str = format!( r#" enable_telemetry = true @@ -2385,6 +3307,8 @@ flight_compression = "arrow_ipc" runtime_size = 8 http2_keep_alive_interval = "10s" http2_keep_alive_timeout = "3s" +enable_cors = false +cors_allowed_origins = [] [grpc.tls] mode = "disable" @@ -2432,7 +3356,7 @@ protocols = [] pending_rows_flush_interval = "0s" max_batch_rows = 100000 max_concurrent_flushes = 256 -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 max_inflight_requests = 3000 flow_notification_queue_capacity = 1024 @@ -2441,7 +3365,6 @@ enable = true [otlp] enable = true -experimental_enable_exponential_histogram = false trace_ingest_chunk_size = 512 experimental_enable_resource_info = true @@ -2449,11 +3372,10 @@ experimental_enable_resource_info = true enable = true with_metric_engine = true prom_validation_mode = "strict" -experimental_enable_prometheus_native_histogram = false pending_rows_flush_interval = "0s" max_batch_rows = 100000 max_concurrent_flushes = 256 -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 max_inflight_requests = 3000 flow_notification_queue_capacity = 1024 @@ -2512,6 +3434,7 @@ experimental_manifest_keep_removed_file_count = 256 experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false experimental_enable_series_index = false +experimental_series_index_max_size = "5GiB" experimental_enable_range_index = false experimental_series_index_maintenance_interval = "5m" experimental_series_index_bucket_width = "5days" @@ -2533,7 +3456,7 @@ scan_memory_on_exhausted = "fail" min_compaction_interval = "0s" schedule_compaction_after_edit = true default_flat_format = true -experimental_series_scan_v2 = true +experimental_series_scan_v2 = false [region_engine.mito.index] aux_path = "" @@ -2561,7 +3484,8 @@ create_on_flush = "auto" create_on_compaction = "auto" apply_on_query = "auto" mem_threshold_on_create = "auto" -{vector_index_config}[region_engine.mito.gc] + +[region_engine.mito.gc] enable = false lingering_time = "1h" unknown_file_lingering_time = "1day" @@ -2618,6 +3542,7 @@ fn drop_lines_with_inconsistent_results(input: String) -> String { "disable_ec2_metadata =", "cache_path =", "cache_capacity =", + "trace_aux_cache_size =", "memory_pool_size =", "scan_memory_limit =", "sas_token =", @@ -2911,6 +3836,18 @@ pub async fn test_prometheus_remote_write(store_type: StorageType) { ) .await; + // A metric-engine logical table stores its rows in the shared physical + // table, so reading them back from SST files goes through a different + // column mapping than the memtable read does. + flush_table(&client, "metric2").await; + validate_data( + "prometheus_remote_write_after_flush", + &client, + "select * from metric2", + table_val, + ) + .await; + // Prom RW tables carry the metric identity; type is inferred from naming. validate_data( "prometheus_remote_write_semantic_identity", @@ -3089,7 +4026,7 @@ pub async fn test_prometheus_remote_write_v2(store_type: StorageType) { pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: StorageType) { common_telemetry::init_default_ut_logging(); - let (app, mut guard) = setup_test_prom_app_with_frontend_native_histogram( + let (app, mut guard) = setup_test_prom_app_with_frontend( store_type, "prometheus_remote_write_v2_native_histogram", ) @@ -3217,11 +4154,24 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ) .await; + let histogram_sql = "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;"; + let histogram_rows = "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]"; validate_data( "prometheus_remote_write_v2_native_histogram_rows", &client, - "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;", - "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", + histogram_sql, + histogram_rows, + ) + .await; + + // The native histogram column is a struct of nested lists, so its SST + // encoding differs from its memtable representation. + flush_table(&client, "remote_write_v2_latency_seconds").await; + validate_data( + "prometheus_remote_write_v2_native_histogram_rows_after_flush", + &client, + histogram_sql, + histogram_rows, ) .await; @@ -3456,7 +4406,7 @@ async fn check_prometheus_remote_write_batched_skip_wal(distributed: bool, v2: b 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 server = build_test_prom_server(instance.frontend(), true, None).build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await; write_prometheus_skip_wal_sample(&client, v2, 1000, None).await; @@ -3530,10 +4480,501 @@ async fn write_prometheus_skip_wal_sample( assert_eq!(request.send().await.status(), StatusCode::NO_CONTENT); } +/// Batched remote write against a logical table bound to a non-millisecond +/// physical table: the bulk guard must reject the destination (the bulk +/// encode only produces millisecond batches) and fall back to the ordinary +/// insert path, which converts the requests to the table's unit. +pub async fn test_prometheus_remote_write_batched_mixed_time_index_units(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_prom_app_with_frontend_batched(store_type, "prom_rw_batched_mixed_units").await; + let client = TestClient::new(app).await; + + let res = client + .get("/v1/sql?db=public&sql=CREATE TABLE phy_us (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, TIME INDEX (greptime_timestamp)) ENGINE = metric WITH ('physical_metric_table' = 'true')") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let write = |metric: &str, value: f64, timestamp: i64| { + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: metric.to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { value, timestamp }], + ..Default::default() + }], + ..Default::default() + }; + prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap() + }; + + // Create the logical table on the microsecond physical table. The bulk + // path handles the microsecond selected physical table: the new table + // is created on it and the samples are widened to its unit during + // batch alignment. + let res = client + .post("/v1/prometheus/write?physical_table=phy_us") + .header("Content-Encoding", "snappy") + .body(write("us_metric", 2.5, 1500)) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + // Write the same existing table again while selecting the default + // (millisecond) physical table: the destination is bound to another + // physical table, so the bulk eligibility check must reject it and fall + // back to the ordinary insert path — without the destination binding + // check the bulk flush would write the rows through the selected + // physical's regions, silently misplacing them. + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(write("us_metric", 3.5, 2000)) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + // Both samples are stored on the microsecond time index. + validate_data( + "prom_rw_batched_mixed_units", + &client, + "SELECT COUNT(*), MAX(greptime_value) FROM us_metric", + "[[2,3.5]]", + ) + .await; + + guard.remove_all().await; +} + +/// A Prometheus write interceptor that redirects every remote write to the +/// `tenant_redirect` schema, like a per-tenant redirection: the incoming +/// context targets `public` while the destination tables live elsewhere. +struct PromSchemaRedirectInterceptor; + +impl servers::interceptor::PromStoreProtocolInterceptor for PromSchemaRedirectInterceptor { + type Error = servers::error::Error; + + fn pre_write( + &self, + _write_req: &RowInsertRequests, + ctx: QueryContextRef, + ) -> servers::error::Result<()> { + ctx.set_current_schema("tenant_redirect"); + Ok(()) + } +} + +async fn setup_redirecting_batched_prom_app( + store_type: StorageType, + name: &str, +) -> (TestClient, TestGuard) { + let plugins = Plugins::default(); + plugins.insert::>( + Arc::new(PromSchemaRedirectInterceptor), + ); + let standalone = tests_integration::standalone::GreptimeDbStandaloneBuilder::new(name) + .with_default_store_type(store_type) + .with_plugin(plugins) + .build() + .await; + let server = build_test_prom_server(standalone.fe_instance().clone(), true, None) + .with_greptime_config_options(standalone.opts.datanode_options().to_toml().unwrap()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + (client, standalone.guard) +} + +async fn create_redirect_schema_with_microsecond_table(client: &TestClient, table: &str) { + for sql in [ + "create database if not exists tenant_redirect", + "CREATE TABLE tenant_redirect.phy_us \ + (ts timestamp(6) time index, val double, host string primary key) \ + engine=metric with ('physical_metric_table' = 'true')", + &format!( + "CREATE TABLE tenant_redirect.{table} \ + (ts timestamp(6) time index, val double, host string primary key) \ + engine=metric with ('on_physical_table' = 'phy_us')" + ), + ] { + let res = client + .get(format!("/v1/sql?sql={sql}").as_str()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "setup: {sql}"); + } +} + +/// Regression test: with batching enabled and a pre_write interceptor that +/// redirects the context to a schema holding same-named non-millisecond +/// tables, the bulk eligibility must be evaluated against the redirected +/// destinations (after preflight) and fall back to the ordinary insert path +/// instead of failing in the bulk encode. +pub async fn test_prometheus_remote_write_batched_interceptor_time_index_units( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (client, mut guard) = + setup_redirecting_batched_prom_app(store_type, "prom_rw_batched_interceptor_units").await; + create_redirect_schema_with_microsecond_table(&client, "intercept_metric").await; + + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "intercept_metric".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prom_rw_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.intercept_metric", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + +/// The v2 variant of the interceptor regression test: a mixed samples-plus- +/// histograms request whose existing sample destination lives in the +/// redirected schema with a microsecond time index must fall back to the +/// ordinary insert path (both series land), not fail in the bulk encode. +pub async fn test_prometheus_remote_write_v2_batched_interceptor_time_index_units( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (client, mut guard) = + setup_redirecting_batched_prom_app(store_type, "prom_rw_v2_batched_interceptor_units") + .await; + create_redirect_schema_with_microsecond_table(&client, "v2_mixed_sample").await; + + // One sample series targeting the existing microsecond table, plus one + // histogram series for a new (millisecond) table. + let mut symbols = vec![String::new()]; + let mut symbol = |value: &str| { + symbols.push(value.to_string()); + (symbols.len() - 1) as u32 + }; + let name_ref = symbol("__name__"); + let sample_ref = symbol("v2_mixed_sample"); + let job_ref = symbol("job"); + let demo_ref = symbol("demo"); + let name_ref2 = symbol("__name__"); + let histo_ref = symbol("v2_mixed_histo"); + let request = RemoteWriteV2Request { + symbols, + timeseries: vec![ + RemoteWriteV2TimeSeries { + labels_refs: vec![name_ref, sample_ref, job_ref, demo_ref], + samples: vec![RemoteWriteV2Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ..Default::default() + }, + RemoteWriteV2TimeSeries { + labels_refs: vec![name_ref2, histo_ref, job_ref, demo_ref], + histograms: vec![Histogram { + count: Some(Count::CountInt(1)), + sum: 1.0, + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_deltas: vec![1], + timestamp: 1000, + start_timestamp: 1000, + ..Default::default() + }], + ..Default::default() + }, + ], + }; + let compressed = prom_store::snappy_compress(&request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prom_rw_v2_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.v2_mixed_sample", + "[[1]]", + ) + .await; + validate_data( + "prom_rw_v2_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.v2_mixed_histo", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + /// 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 /// shared request context. +/// Regression test for : +/// prometheus remote write uses the logical batcher against a physical +/// metric table pre-created with a microsecond time index; the millisecond +/// samples are widened to the physical table's unit during batch alignment. +pub async fn test_prometheus_remote_write_batched_microsecond_physical_table( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_prom_app_with_frontend_batched( + store_type, + "prometheus_remote_write_batched_us_physical", + ) + .await; + let client = TestClient::new(app).await; + + // Pre-create the default physical metric table with a microsecond time + // index before any remote write, so the batched bulk path must handle it. + let res = client + .get( + "/v1/sql?db=public&sql=CREATE TABLE greptime_physical_table \ + (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')", + ) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let write_request = WriteRequest { + timeseries: vec![prom_store::mock_timeseries()[0].clone()], + ..Default::default() + }; + let serialized_request = write_request.encode_to_vec(); + let compressed_request = + prom_store::snappy_compress(&serialized_request).expect("failed to encode snappy"); + + // submit_build_and_align counts every batcher submission in both + // acknowledgement modes, so the stored-value assertions below cannot be + // satisfied by a silent fallback to ordinary insertion. + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + let before = submissions.get_sample_count(); + + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed_request) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 1, + "non-millisecond physical tables must use the logical batcher" + ); + + // metric1 samples are 1.0@1000ms and 2.0@2000ms; on the microsecond + // physical table they must be stored as 1_000_000us and 2_000_000us. + wait_for_data( + &client, + "select greptime_timestamp, greptime_value from metric1 order by greptime_timestamp", + "[[1000000,1.0],[2000000,2.0]]", + ) + .await; + + guard.remove_all().await; +} + +/// Regression test for the per-series physical-table selection conflict: +/// one remote-write request containing the same metric under two different +/// physical-table selections must fall back to the ordinary insert path — +/// both when the destination table already exists (bound to one of the two +/// physicals) and when it is missing (conflicting creations). A consistent +/// selection must still use the batcher. +pub async fn test_prometheus_remote_write_batched_conflicting_physical_selections( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_prom_app_with_frontend_batched( + store_type, + "prom_rw_batched_conflicting_physicals", + ) + .await; + let client = TestClient::new(app).await; + + let res = client + .get( + "/v1/sql?db=public&sql=CREATE TABLE p1 (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, TIME INDEX (greptime_timestamp)) ENGINE = metric WITH ('physical_metric_table' = 'true')", + ) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + // An existing logical table bound to p1. + let res = client + .get( + "/v1/sql?db=public&sql=CREATE TABLE conflict_existing (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \"job\" STRING NULL, TIME INDEX (greptime_timestamp), PRIMARY KEY (\"job\")) ENGINE = metric WITH ('on_physical_table' = 'p1')", + ) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let write = |metric: &str, job: &str, physical: Option<&str>, value: f64| { + let mut labels = vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: metric.to_string(), + }, + Label { + name: "job".to_string(), + value: job.to_string(), + }, + ]; + if let Some(physical) = physical { + labels.push(Label { + name: "x_greptime_physical_table".to_string(), + value: physical.to_string(), + }); + } + TimeSeries { + labels, + samples: vec![Sample { + value, + timestamp: 1000, + }], + ..Default::default() + } + }; + let send = |timeseries: Vec| { + let write_request = WriteRequest { + timeseries, + ..Default::default() + }; + prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap() + }; + + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + + // Existing destination under two different selections: the batch must + // fall back (the batch for the unbound selection would otherwise write + // through the default physical's regions), and both series land on the + // bound physical. + let before = submissions.get_sample_count(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(send(vec![ + write("conflict_existing", "a", Some("p1"), 1.0), + write("conflict_existing", "b", None, 2.0), + ])) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 0, + "conflicting selections for an existing table must fall back" + ); + wait_for_data( + &client, + "select count(*), sum(greptime_value) from conflict_existing", + "[[2,3.0]]", + ) + .await; + + // Missing destination under two different selections: conflicting + // creations must also fall back, and both series must be visible + // through the created table. + let before = submissions.get_sample_count(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(send(vec![ + write("conflict_fresh", "a", Some("p1"), 1.0), + write("conflict_fresh", "b", None, 2.0), + ])) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 0, + "conflicting selections for a missing table must fall back" + ); + wait_for_data( + &client, + "select count(*), sum(greptime_value) from conflict_fresh", + "[[2,3.0]]", + ) + .await; + + // A consistent selection still uses the batcher. + let before = submissions.get_sample_count(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(send(vec![ + write("consistent_fresh", "a", Some("p1"), 1.0), + write("consistent_fresh", "b", Some("p1"), 2.0), + ])) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 1, + "a consistent selection must use the logical batcher" + ); + wait_for_data( + &client, + "select count(*), sum(greptime_value) from consistent_fresh", + "[[2,3.0]]", + ) + .await; + + guard.remove_all().await; +} + pub async fn test_prometheus_remote_write_batched(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -4283,6 +5724,31 @@ transform: ) .await; + // The inverted, skipping and fulltext index files are only written on flush, + // so the queries below are the first thing that reads them back. + flush_table(&client, "test_db.logs1").await; + validate_data( + "pipeline_db_inverted_index", + &client, + "select id2 from test_db.logs1 where id1 = 2436", + "[[2528]]", + ) + .await; + validate_data( + "pipeline_db_skipping_index", + &client, + "select id1 from test_db.logs1 where type = 'I'", + "[[2436]]", + ) + .await; + validate_data( + "pipeline_db_fulltext_index", + &client, + "select id1 from test_db.logs1 where matches(log, 'ClusterAdapter')", + "[[2436]]", + ) + .await; + // 5. remove pipeline let encoded_ver_str: String = url::form_urlencoded::byte_serialize(version_str.as_bytes()).collect(); @@ -5079,6 +6545,24 @@ transform: ) .await; + // The fulltext and skipping index files are only written on flush, so the + // queries below are the first thing that reads them back. + flush_table(&client, "pipeline_index_options").await; + validate_data( + "pipeline_index_options_fulltext", + &client, + "select trace_id from pipeline_index_options where matches(message, 'greptime')", + "[[42]]", + ) + .await; + validate_data( + "pipeline_index_options_skipping", + &client, + "select message from pipeline_index_options where trace_id = 42", + r#"[["hello greptime"]]"#, + ) + .await; + guard.remove_all().await; } @@ -6995,6 +8479,18 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // A metric-engine logical table stores its rows in the shared physical + // table, so reading them back from SST files goes through a different + // column mapping than the memtable read does. + flush_table(&client, "claude_code_cost_usage_USD_total").await; + validate_data( + "otlp_metrics_all_select_after_flush", + &client, + "select * from `claude_code_cost_usage_USD_total` order by model desc;", + expected, + ) + .await; + // The synthesized resource descriptor: a plain mito info table keyed by // the raw OTel attribute keys (service.name is the only allowlisted // resource attribute in this payload), independent of the promote/ignore @@ -7333,7 +8829,6 @@ pub async fn test_otlp_exponential_histogram(store_type: StorageType) { AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Metric, ResourceMetrics, ScopeMetrics, exponential_histogram_data_point, metric, }; - use tests_integration::test_util::setup_test_http_app_with_otlp_exponential_histogram; common_telemetry::init_default_ut_logging(); let req = ExportMetricsServiceRequest { @@ -7380,44 +8875,8 @@ pub async fn test_otlp_exponential_histogram(store_type: StorageType) { )] }; - let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( - store_type, - "test_otlp_exponential_histogram_disabled", - false, - ) - .await; - let client = TestClient::new(app).await; - let res = send_req( - &client, - headers(), - "/v1/otlp/v1/metrics", - body.clone(), - false, - ) - .await; - assert_eq!(StatusCode::BAD_REQUEST, res.status()); - let status = GoogleRpcStatus::decode(res.bytes().await.as_ref()).unwrap(); - assert_eq!(3, status.code); - assert!( - status - .message - .contains("otlp.experimental_enable_exponential_histogram") - ); - validate_data( - "otlp_exponential_histogram_disabled_no_table", - &client, - "select count(*) from information_schema.tables where table_name = 'otlp_exponential_latency';", - "[[0]]", - ) - .await; - guard.remove_all().await; - - let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( - store_type, - "test_otlp_exponential_histogram_enabled", - true, - ) - .await; + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_exponential_histogram").await; let client = TestClient::new(app).await; let res = send_req(&client, headers(), "/v1/otlp/v1/metrics", body, false).await; assert_eq!(StatusCode::OK, res.status()); @@ -7846,6 +9305,17 @@ pub async fn test_otlp_traces_v0(store_type: StorageType) { ) .await; + // The attribute columns are JSON and the event/link columns are lists, whose + // SST encodings differ from their memtable representations. + flush_table(&client, "opentelemetry_traces").await; + validate_data( + "otlp_traces_after_flush", + &client, + "select * from opentelemetry_traces;", + expected, + ) + .await; + guard.remove_all().await; } @@ -7972,6 +9442,30 @@ pub(crate) async fn test_otlp_traces_v2( ) .await; + // JSON2 path extraction has to resolve the same way against SST files as it + // does against the memtable. + flush_table(&client, "trace_v2_spans").await; + validate_data( + "otlp_traces_v2_json2_after_flush", + &client, + "select service_name, span_attributes.\"http.status_code\"::BIGINT, \ + resource_attributes.\"deployment.environment\"::STRING, \ + scope_attributes.enabled::BOOLEAN from trace_v2_spans;", + r#"[["frontend",200,"production",true]]"#, + ) + .await; + validate_data( + "otlp_traces_v2_events_links_after_flush", + &client, + r#"select json_get_string(span_events, '$[0].name'), + json_get_int(span_events, '$[0].attributes."event.code"'), + json_get_string(span_links, '$[0].trace_id'), + json_get_string(span_links, '$[0].attributes."link.type"') + from trace_v2_spans;"#, + r#"[["cache.hit",7,"cc9e0991a2e63d274984bd44ee669203","follows_from"]]"#, + ) + .await; + guard.remove_all().await; Ok(()) } @@ -8125,6 +9619,17 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { let expected = r#"[[1736480942444376000,1736480942444499000,123000,null,"c05d7a4ec8e1f231f02ed6e8da8655b4","d24f921c75f68e23","SPAN_KIND_CLIENT","lets-go","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-server",[],[]],[1736480942444376000,1736480942444499000,123000,"d24f921c75f68e23","c05d7a4ec8e1f231f02ed6e8da8655b4","9630f2916e2f7909","SPAN_KIND_SERVER","okey-dokey-0","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-client",[],[]],[1736480942444589000,1736480942444712000,123000,null,"cc9e0991a2e63d274984bd44ee669203","eba7be77e3558179","SPAN_KIND_CLIENT","lets-go","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-server",[],[]],[1736480942444589000,1736480942444712000,123000,"eba7be77e3558179","cc9e0991a2e63d274984bd44ee669203","8f847259b0f6e1ab","SPAN_KIND_SERVER","okey-dokey-0","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-client",[],[]]]"#; validate_data("otlp_traces", &client, "select * from mytable;", expected).await; + // `span_events` / `span_links` are JSON columns and the table is partitioned + // on `trace_id`, so the read back has to reassemble SST files across regions. + flush_table(&client, "mytable").await; + validate_data( + "otlp_traces_after_flush", + &client, + "select * from mytable;", + expected, + ) + .await; + // The trace v1 main table carries the trace identity (events/links preserved as // JSON columns by the v1 model). validate_data( @@ -8961,6 +10466,17 @@ pub async fn test_otlp_logs(store_type: StorageType) { ) .await; + // The attribute columns are JSON, whose SST encoding differs from their + // memtable representation. + flush_table(&client, "opentelemetry_logs").await; + validate_data( + "otlp_logs_after_flush", + &client, + "select * from opentelemetry_logs;", + expected, + ) + .await; + // The auto-created log table carries the log identity. validate_data( "otlp_logs_semantic_identity", @@ -9343,6 +10859,17 @@ pub async fn test_loki_pb_logs(store_type: StorageType) { ) .await; + // `structured_metadata` is a JSON column, whose SST encoding differs from its + // memtable representation. + flush_table(&client, "loki_table_name").await; + validate_data( + "loki_pb_content_after_flush", + &client, + "select * from loki_table_name;", + expected, + ) + .await; + guard.remove_all().await; } @@ -9548,6 +11075,17 @@ pub async fn test_loki_json_logs(store_type: StorageType) { ) .await; + // `structured_metadata` is a JSON column, whose SST encoding differs from its + // memtable representation. + flush_table(&client, "loki_table_name").await; + validate_data( + "loki_json_content_after_flush", + &client, + "select * from loki_table_name;", + expected, + ) + .await; + guard.remove_all().await; } @@ -9767,7 +11305,7 @@ pub async fn test_log_query(store_type: StorageType) { // prepare data with SQL API let res = client - .get("/v1/sql?sql=create table logs (`ts` timestamp time index, `message` string);") + .get("/v1/sql?sql=create table logs (`ts` timestamp time index, `message` string fulltext index with(backend='bloom'));") .send() .await; assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); @@ -9842,6 +11380,35 @@ pub async fn test_log_query(store_type: StorageType) { ] ); + // The fulltext index file is only written on flush. Replay the log query to + // cover the SST read path, and add a `matches_term` query, which is the + // predicate the index applier actually recognises. + flush_table(&client, "logs").await; + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&log_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + let resp = res.text().await; + let output = serde_json::from_str::(&resp).unwrap(); + let rows = output["output"][0]["records"]["rows"].as_array().unwrap(); + let mut flushed_messages = rows + .iter() + .map(|row| row[1].as_str().unwrap()) + .collect::>(); + flushed_messages.sort_unstable(); + assert_eq!(flushed_messages, messages); + + validate_data( + "log_query_fulltext_index", + &client, + "select message from logs where matches_term(message, 'at-explicit-end')", + r#"[["at-explicit-end"]]"#, + ) + .await; + let res = client .get("/v1/sql?sql=create table logs_limit (`ts` timestamp time index, `message` string);") .send() @@ -10865,11 +12432,20 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let mut req: ExportTraceServiceRequest = serde_json::from_str(content).unwrap(); - // Modify timestamp fields + // The fixture is written against a fixed base instant. Move it next to now: + // this table is created with `ttl=7d`, and stale spans would be dropped as + // expired the moment they are flushed into SST files. let now = Utc::now().timestamp_nanos_opt().unwrap() as u64; + let delta_micros = (now / 1_000) as i64 - 60_000_000 - JAEGER_V1_FIXTURE_BASE_MICROS as i64; + let delta_nanos = delta_micros * 1_000; for span in req.resource_spans.iter_mut() { for scope_span in span.scope_spans.iter_mut() { - // Only modify the timestamp fields for the span with the name "test-jaeger-get-operations" to current time. + for span in scope_span.spans.iter_mut() { + span.start_time_unix_nano = (span.start_time_unix_nano as i64 + delta_nanos) as u64; + span.end_time_unix_nano = (span.end_time_unix_nano as i64 + delta_nanos) as u64; + } + // `/api/operations` is queried without a time range, so these spans + // stay pinned to the last few seconds. if scope_span.scope.as_ref().unwrap().name == "test-jaeger-get-operations" { for span in scope_span.spans.iter_mut() { span.start_time_unix_nano = now - 5_000_000_000; // 5 seconds ago @@ -10946,7 +12522,8 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/operations` API. @@ -10995,12 +12572,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/services/{service_name}/operations` API. let res = client - .get("/v1/jaeger/api/services/test-jaeger-query-api/operations?start=1738726754492421&end=1738726754642422") + .get(&shift_fixture_micros("/v1/jaeger/api/services/test-jaeger-query-api/operations?start=1738726754492421&end=1738726754642422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11023,7 +12601,8 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API without start and end. @@ -11149,12 +12728,26 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); + assert_eq!(resp, expected); + + // Replay the same lookup against SST files: `trace_id` and `service_name` + // carry BLOOM skipping indexes whose puffin files only exist after a flush, + // and `span_events` / `span_links` are JSON columns. + flush_table(&client, trace_table_name).await; + let res = client + .get("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea") + .header("x-greptime-trace-table-name", trace_table_name) + .send() + .await; + assert_eq!(StatusCode::OK, res.status()); + let resp: Value = serde_json::from_str(&res.text().await).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API with start and end in microseconds. let res = client - .get("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea?start=1738726754492421&end=1738726754642422") + .get(&shift_fixture_micros("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea?start=1738726754492421&end=1738726754642422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11275,12 +12868,16 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API for non-existent trace. let res = client - .get("/v1/jaeger/api/traces/0000000000000000000000000000dead") + .get(&shift_fixture_micros( + "/v1/jaeger/api/traces/0000000000000000000000000000dead", + delta_micros, + )) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11299,12 +12896,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&operation=access-mysql&start=1738726754492421&end=1738726754642422&tags=%7B%22operation.type%22%3A%22access-mysql%22%7D") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&operation=access-mysql&start=1738726754492421&end=1738726754642422&tags=%7B%22operation.type%22%3A%22access-mysql%22%7D", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11426,13 +13024,14 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API with tags. // 1. first query without tags, get 2 results let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11443,12 +13042,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // 2. second query with tags, get 1 result let res = client -.get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422&tags=%7B%22operation.type%22%3A%22access-pg%22%7D") +.get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422&tags=%7B%22operation.type%22%3A%22access-pg%22%7D", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11459,14 +13059,15 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API with Grafana User-Agent. // When user agent is Grafana, only return at most 3 spans per trace (earliest by timestamp). // Trace `5611dce1bc9ebed65352d99a027b08fb` has 4 spans, so only 3 should be returned. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .header("User-Agent", "Grafana/8.0.0") .send() @@ -11514,7 +13115,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { // Test `/api/traces` API without User-Agent (default behavior). // All 4 spans should be returned for the trace. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11534,7 +13135,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { // Test `/api/traces` API with Jaeger User-Agent (should return all spans like default). let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .header("User-Agent", "Jaeger-Query/1.0.0") .send() @@ -11613,6 +13214,59 @@ pub async fn test_influxdb_write(store_type: StorageType) { guard.remove_all().await; } +/// Flushes `table` so the rows written so far move from the memtable into SST +/// files. Replaying a query afterwards covers the encode/decode round trip and, +/// for indexed columns, the puffin files that only exist after a flush. +/// Base microsecond instant the `test_jaeger_query_api_for_trace_v1` fixture and +/// its expected responses are written against. +const JAEGER_V1_FIXTURE_BASE_MICROS: u64 = 1_738_726_754_492_421; + +/// Shifts every fixture timestamp in `text` by `delta` microseconds, so expected +/// responses and query bounds follow the payload after it is moved next to now. +/// +/// Only integer literals within a day of [`JAEGER_V1_FIXTURE_BASE_MICROS`] are +/// rewritten; durations, counts and ids are left alone. +fn shift_fixture_micros(text: &str, delta: i64) -> String { + let mut out = String::with_capacity(text.len()); + let mut rest = text; + while let Some(start) = rest.find(|c: char| c.is_ascii_digit()) { + out.push_str(&rest[..start]); + let digits_end = rest[start..] + .find(|c: char| !c.is_ascii_digit()) + .map_or(rest.len(), |offset| start + offset); + let digits = &rest[start..digits_end]; + match digits.parse::() { + Ok(value) if (value - JAEGER_V1_FIXTURE_BASE_MICROS as i64).abs() < 86_400_000_000 => { + out.push_str(&(value + delta).to_string()) + } + _ => out.push_str(digits), + } + rest = &rest[digits_end..]; + } + out.push_str(rest); + out +} + +async fn flush_table(client: &TestClient, table: &str) { + // The argument is parsed as a SQL table name, so an unquoted identifier is + // lower-cased and will not resolve a mixed-case table. + let quoted = table + .split('.') + .map(|part| format!("\"{part}\"")) + .collect::>() + .join("."); + let res = client + .post("/v1/sql") + .form(&[("sql", format!("ADMIN FLUSH_TABLE('{quoted}')"))]) + .send() + .await; + let status = res.status(); + let body = res.text().await; + assert_eq!(status, StatusCode::OK, "flush {table} fail: {body}"); + // A failing statement still comes back as 200 with an `error` envelope. + assert!(!body.contains("\"error\""), "flush {table} fail: {body}"); +} + async fn validate_data(test_name: &str, client: &TestClient, sql: &str, expected: &str) { let res = client .get(format!("/v1/sql?sql={sql}").as_str()) @@ -11715,16 +13369,9 @@ async fn check_http_skip_wal(name: &str, cases: &[HttpWalCase], distributed: boo .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) + .with_otlp_handler(fe.clone(), true) // The pending batcher uses BulkInsert, deliberately outside this PR. - .with_prom_handler( - fe.clone(), - Some(fe), - true, - PromValidationMode::Strict, - false, - None, - ) + .with_prom_handler(fe.clone(), Some(fe), true, PromValidationMode::Strict, None) .build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await; for case in cases { @@ -12069,7 +13716,6 @@ pub async fn test_http_memory_limit(store_type: StorageType) { None, Some(http_opts), Some(memory_limiter), - false, ) .await; diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index b61b6cd3e19..72d0a1850bd 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -15,6 +15,7 @@ #![recursion_limit = "256"] mod admin_function_event; +mod build_series_index; mod database_ddl_event; mod dict_groupby_sst; mod event_recorder_test_util; @@ -38,6 +39,7 @@ mod repartition_event; mod repartition_expr_version; mod export_logical_tables; mod mysql; +mod object_store_wal; mod reconciliation_event; mod view_ddl_event; mod wal_prune_event; @@ -53,6 +55,17 @@ grpc_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); http_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); +http_local_tests!( + test_health_api, + test_status_api, + test_metrics_api, + test_splunk_health, + test_splunk_health_is_public, + test_dynamic_tracer_toggle, + test_dashboard_path, + test_test_pipeline_api, +); + #[tokio::test(flavor = "multi_thread")] async fn test_http_auth_from_standalone_user_provider_config() { http::test_http_auth_from_standalone_user_provider_config().await; diff --git a/tests-integration/tests/object_store_wal.rs b/tests-integration/tests/object_store_wal.rs new file mode 100644 index 00000000000..0aef1e145b0 --- /dev/null +++ b/tests-integration/tests/object_store_wal.rs @@ -0,0 +1,211 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::path::Path; +use std::time::Duration; + +use common_procedure::options::ProcedureConfig; +use common_query::Output; +use common_wal::config::DatanodeWalConfig; +use common_wal::config::object_store::ObjectStoreWalConfig; +use frontend::instance::Instance; +use object_store::ObjectStore; +use object_store::config::ObjectStoreConfig; +use object_store::services::{Fs, S3}; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use tests_integration::standalone::{GreptimeDbStandalone, GreptimeDbStandaloneBuilder}; +use tests_integration::test_util::StorageType; + +/// The root the store derives its node prefix from. +const WAL_ROOT: &str = "cluster-a/wal"; +const BATCHES: usize = 5; +const ROWS_PER_BATCH: usize = 4; +const QUERY: &str = "SELECT hostname, usage_user, ts FROM cpu ORDER BY ts"; + +async fn execute_sql(instance: &Instance, sql: &str) -> Output { + SqlQueryHandler::do_query(instance, sql, QueryContext::arc()) + .await + .remove(0) + .unwrap() +} + +async fn query_rows(standalone: &GreptimeDbStandalone) -> String { + execute_sql(standalone.fe_instance(), QUERY) + .await + .data + .pretty_print() + .await +} + +/// Returns the number of WAL objects under [`WAL_ROOT`] of the default store. +async fn count_wal_objects(standalone: &GreptimeDbStandalone) -> usize { + let store = match &standalone.opts.storage.store { + ObjectStoreConfig::File(_) => { + ObjectStore::new(Fs::default().root(&standalone.opts.storage.data_home)).unwrap() + } + ObjectStoreConfig::S3(s3) => ObjectStore::new(S3::from(&s3.connection)).unwrap(), + other => panic!("unexpected store {other:?}"), + }; + store + .list_with(&format!("{WAL_ROOT}/")) + .recursive(true) + .await + .unwrap() + .into_iter() + .filter(|entry| entry.metadata().is_file()) + .count() +} + +/// Waits until collection leaves `expected` WAL objects under [`WAL_ROOT`]. +async fn wait_for_wal_objects(standalone: &GreptimeDbStandalone, expected: usize) { + let collected = async { + while count_wal_objects(standalone).await != expected { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }; + tokio::time::timeout(Duration::from_secs(30), collected) + .await + .unwrap_or_else(|_| panic!("expected {expected} WAL objects after collection")); +} + +/// Stops the procedure manager, the region server and the WAL of the +/// instance. +async fn shutdown(standalone: &mut GreptimeDbStandalone) { + standalone.procedure_manager.stop().await.unwrap(); + standalone.datanode.shutdown().await.unwrap(); +} + +/// Shuts the instance down and builds it again on its metadata, data home and +/// object store, like a process restart. +async fn restart( + builder: &GreptimeDbStandaloneBuilder, + mut standalone: GreptimeDbStandalone, +) -> GreptimeDbStandalone { + shutdown(&mut standalone).await; + let GreptimeDbStandalone { + opts, + guard, + kv_backend, + .. + } = standalone; + + let (procedure_manager, event_recorder_handle) = + standalone::build_procedure_manager(kv_backend.clone(), ProcedureConfig::default()); + builder + .build_with( + kv_backend, + guard, + opts, + procedure_manager, + event_recorder_handle, + true, + ) + .await +} + +/// Writes batches that land in separate WAL objects, restarts before and +/// after a flush, and checks that the rows survive both restarts and that the +/// flush lets collection delete every earlier object. +async fn run_restarts_on_object_store_wal(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let wal_config = ObjectStoreWalConfig { + prefix: WAL_ROOT.to_string(), + ..Default::default() + }; + let builder = GreptimeDbStandaloneBuilder::new("object_store_wal") + .with_default_store_type(store_type) + .with_datanode_wal_config(DatanodeWalConfig::ObjectStore(wal_config)); + let standalone = builder.build().await; + + execute_sql( + standalone.fe_instance(), + r#" + CREATE TABLE cpu ( + hostname STRING PRIMARY KEY, + usage_user DOUBLE, + ts TIMESTAMP TIME INDEX + ) + "#, + ) + .await; + // Every insert returns once its entries are durable, so the batches land + // in separate WAL objects after the object that started the epoch. + for batch in 0..BATCHES { + insert_batch(&standalone, batch).await; + } + let expected = query_rows(&standalone).await; + assert_eq!(BATCHES * ROWS_PER_BATCH + 4, expected.lines().count()); + // The region has no watermark, so its first object holds the boundary + // and only the object that started the epoch is collected. + wait_for_wal_objects(&standalone, BATCHES).await; + let objects = BATCHES; + // No Raft Engine log store is created as a fallback. + assert!( + !Path::new(&standalone.opts.storage.data_home) + .join("wal") + .exists() + ); + + // Nothing was flushed, so the rows come back from the WAL alone; the + // reopened store starts its epoch with another object. + let standalone = restart(&builder, standalone).await; + assert_eq!(expected, query_rows(&standalone).await); + assert!(count_wal_objects(&standalone).await > objects); + + // After the flush the next object carries a reclaim boundary past every + // earlier object, so collection leaves that object alone. The flushed + // rows come back from the SST files and the last batch from the WAL. + execute_sql(standalone.fe_instance(), "ADMIN FLUSH_TABLE('cpu')").await; + insert_batch(&standalone, BATCHES).await; + wait_for_wal_objects(&standalone, 1).await; + let expected = query_rows(&standalone).await; + let mut standalone = restart(&builder, standalone).await; + assert_eq!(expected, query_rows(&standalone).await); + shutdown(&mut standalone).await; +} + +/// Inserts the rows of `batch`, which land in one WAL object. +async fn insert_batch(standalone: &GreptimeDbStandalone, batch: usize) { + let values = (0..ROWS_PER_BATCH) + .map(|row| { + let i = batch * ROWS_PER_BATCH + row; + format!( + "('host_{i}', {i}.0, {})", + 1_686_567_600_000 + i as i64 * 1000 + ) + }) + .collect::>() + .join(", "); + execute_sql( + standalone.fe_instance(), + &format!("INSERT INTO cpu VALUES {values}"), + ) + .await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_standalone_object_store_wal_survives_restarts_on_file() { + run_restarts_on_object_store_wal(StorageType::File).await; +} + +/// Runs against the S3 bucket of the `GT_S3_*` environment variables, so it +/// is skipped unless they are set. +#[tokio::test(flavor = "multi_thread")] +async fn test_standalone_object_store_wal_survives_restarts_on_s3() { + if !StorageType::S3.test_on() { + return; + } + run_restarts_on_object_store_wal(StorageType::S3).await; +} diff --git a/tests-integration/tests/reconciliation_event.rs b/tests-integration/tests/reconciliation_event.rs index 19eaea6ff2c..a366217a7d9 100644 --- a/tests-integration/tests/reconciliation_event.rs +++ b/tests-integration/tests/reconciliation_event.rs @@ -20,6 +20,7 @@ use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; use table::table_reference::TableReference; use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; @@ -33,6 +34,7 @@ async fn test_catalog_and_database_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("catalog_database_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; @@ -243,6 +245,7 @@ async fn test_table_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("table_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; @@ -351,6 +354,7 @@ async fn test_logical_table_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("logical_table_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index c243113d928..1198e6d23e5 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -45,15 +45,25 @@ macro_rules! repartition_tests { $( paste::item! { mod [] { + // Every case below builds its own cluster and runs a full repartition plus + // GC cycle, which takes around a minute on object-store backends. Keep one + // case per test so a case stays well inside the nextest slow timeout and + // cases run in parallel; do not fold several cases back into one test. + #[tokio::test(flavor = "multi_thread")] - async fn [< test_repartition_mito >]() { + async fn [< test_repartition_mito_flat >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { common_telemetry::init_default_ut_logging(); - // Cover both storage formats for repartition behavior. - // for flat format $crate::repartition::test_repartition_mito(store_type, true).await; - // for primary key format + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_mito_primary_key >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_mito(store_type, false).await; } } @@ -95,19 +105,41 @@ macro_rules! repartition_tests { } #[tokio::test(flavor = "multi_thread")] - async fn [< test_repartition_metric >]() { + async fn [< test_repartition_metric_flat_sparse >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { use store_api::codec::PrimaryKeyEncoding; common_telemetry::init_default_ut_logging(); - // Exercise format + primary key encoding matrix for metric engine. - // for flat format with sparse primary key encoding $crate::repartition::test_repartition_metric(store_type, true, PrimaryKeyEncoding::Sparse).await; - // for flat format with dense primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_flat_dense >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, true, PrimaryKeyEncoding::Dense).await; - // for primary key format with sparse primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_primary_key_sparse >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, false, PrimaryKeyEncoding::Sparse).await; - // for primary key format with dense primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_primary_key_dense >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, false, PrimaryKeyEncoding::Dense).await; } } @@ -117,6 +149,101 @@ macro_rules! repartition_tests { }; } +#[tokio::test(flavor = "multi_thread")] +async fn test_repartition_physical_metric_without_logical_table() { + common_telemetry::init_default_ut_logging(); + let (store_config, _guard) = get_test_store_config(&StorageType::File); + let home_dir = create_temp_dir("repartition_physical_metric_without_logical_table"); + let cluster = + GreptimeDbClusterBuilder::new("test_repartition_physical_metric_without_logical_table") + .await + .with_datanodes(3) + .with_store_config(store_config) + .with_shared_home_dir(Arc::new(home_dir)) + .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) + .build(true) + .await; + + let instance = cluster.fe_instance(); + let query_ctx = QueryContext::arc(); + let sql = r#" + CREATE TABLE `physical_metric_without_logical`( + `ts` TIMESTAMP TIME INDEX, + `val` DOUBLE, + `host` STRING PRIMARY KEY + ) PARTITION ON COLUMNS (`host`) ( + `host` < 'm', + `host` >= 'm' + ) ENGINE = metric + WITH ( + "physical_metric_table" = "true" + ); + "#; + run_sql(instance, sql, query_ctx.clone()).await.unwrap(); + + let table_id = get_table_id(&cluster.metasrv, "physical_metric_without_logical").await; + let table_info = cluster + .metasrv + .table_metadata_manager() + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + let column_ids_before_split = table_info.table_info.meta.column_ids.clone(); + assert_eq!(column_ids_before_split, vec![0, 1, 2]); + + let sql = r#" + ALTER TABLE `physical_metric_without_logical` SPLIT PARTITION ( + `host` < 'm' + ) INTO ( + `host` < 'g', + `host` >= 'g' AND `host` < 'm' + ); + "#; + run_sql(instance, sql, query_ctx.clone()).await.unwrap(); + + let table_info = cluster + .metasrv + .table_metadata_manager() + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + assert_eq!( + table_info.table_info.meta.column_ids, + column_ids_before_split + ); + + let result = run_sql( + instance, + &query_partitions_sql("physical_metric_without_logical"), + query_ctx, + ) + .await + .unwrap(); + let expected = r#"+---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+ +| table_catalog | table_schema | table_name | partition_name | partition_expression | partition_description | greptime_partition_id | partition_ordinal_position | ++---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+ +| greptime | public | physical_metric_without_logical | p0 | host | host < g | 4398046511104 | 1 | +| greptime | public | physical_metric_without_logical | p1 | host | host >= m | 4398046511105 | 2 | +| greptime | public | physical_metric_without_logical | p2 | host | host >= g AND host < m | 4398046511106 | 3 | ++---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+"#; + check_output_stream(result.data, expected).await; +} + /// COUNT must use visible rows rather than the full row counts of shared SSTs. #[tokio::test(flavor = "multi_thread")] async fn test_repartition_append_count_file() { diff --git a/tests-integration/tests/repartition_event.rs b/tests-integration/tests/repartition_event.rs index 906eb3cfee3..aaefcc2d49c 100644 --- a/tests-integration/tests/repartition_event.rs +++ b/tests-integration/tests/repartition_event.rs @@ -22,6 +22,7 @@ use mito2::gc::GcConfig; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::test_util::{ StorageType, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use crate::event_recorder_test_util::{ @@ -44,6 +45,7 @@ async fn test_repartition_event() { let home_dir = create_temp_dir("test_repartition_event_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_repartition_event") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_store_config(store_config) .with_shared_home_dir(Arc::new(home_dir)) diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 6c58ab0246c..cb3d151fa85 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -31,7 +31,8 @@ use sqlx::{Connection, Executor, Row}; use tests_integration::test_util::{ StorageType, setup_mysql_server, setup_mysql_server_with_slow_query_threshold, setup_mysql_server_with_user_provider, setup_pg_server, - setup_pg_server_with_slow_query_threshold, setup_pg_server_with_user_provider, + setup_pg_server_with_prom_native_histogram, setup_pg_server_with_slow_query_threshold, + setup_pg_server_with_user_provider, }; use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; @@ -94,6 +95,8 @@ macro_rules! sql_tests { test_mysql_federated_prepare_stmt, test_mysql_prepare_tql_and_show, test_postgres_extended_query_row_returning_statements, + test_postgres_native_histogram, + test_postgres_struct_types, test_declare_fetch_close_cursor, test_alter_update_on, ); @@ -879,6 +882,234 @@ pub async fn test_postgres_bytea(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_postgres_native_histogram(store_type: StorageType) { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use axum::http::StatusCode; + use prost::Message; + use servers::http::test_helpers::TestClient; + use servers::prom_remote_write::v2::test_util as remote_write_v2; + use servers::prom_store; + + let (mut guard, app, fe_pg_server) = + setup_pg_server_with_prom_native_histogram(store_type, "test_postgres_native_histogram") + .await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + let http_client = TestClient::new(app).await; + + // Ingest one integer-count and one float-count native histogram via + // prometheus remote-write v2. + let write_request = remote_write_v2::request_with_labels_and_histograms( + vec![ + (prom_store::METRIC_NAME_LABEL, "pg_native_histogram_seconds"), + ("job", "api"), + ("instance", "localhost:9090"), + ], + vec![ + Histogram { + count: Some(Count::CountInt(8)), + sum: 10.0, + schema: 1, + zero_threshold: 0.001, + zero_count: Some(ZeroCount::ZeroCountInt(1)), + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + reset_hint: 2, + timestamp: 3000, + start_timestamp: 1500, + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(6.0)), + sum: 20.0, + schema: 2, + zero_threshold: 0.002, + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + reset_hint: 3, + timestamp: 4000, + start_timestamp: 2500, + ..Default::default() + }, + ], + ); + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = http_client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + let rows = client + .simple_query( + "select greptime_native_histogram from pg_native_histogram_seconds order by greptime_timestamp", + ) + .await + .unwrap(); + let jsons: Vec<&str> = rows + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0), + _ => None, + }) + .collect(); + assert_eq!(jsons.len(), 2); + + // The struct column is encoded as a JSON object, matching the HTTP query + // output for the same histograms. Before the fix every row was encoded as + // an empty struct (`{}`). + let histogram: serde_json::Value = serde_json::from_str(jsons[0]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let histogram: serde_json::Value = serde_json::from_str(jsons[1]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + +pub async fn test_postgres_struct_types(store_type: StorageType) { + let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_postgres_struct_types").await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + async fn query_one(client: &Client, sql: &str) -> String { + let messages = client.simple_query(sql).await.unwrap(); + messages + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0).map(|v| v.to_string()), + _ => None, + }) + .next() + .unwrap() + } + + // Struct and list-of-struct columns are serialized as JSON. View-typed + // struct fields must not fail row extraction, and a null struct inside a + // list stays null instead of becoming a struct of null fields. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + // '0102' casts to its UTF-8 bytes when interpreted as BinaryView. + let row = query_one(&client, "SELECT struct(arrow_cast('0102', 'BinaryView'))").await; + assert_eq!(row, "{\"c0\":[48,49,48,50]}"); + + let row = query_one(&client, "SELECT struct([struct(1), NULL])").await; + assert_eq!(row, "{\"c0\":[{\"c0\":1},null]}"); + + // Top-level control: the same null element in a bare list stays null too. + let row = query_one(&client, "SELECT [struct(1), NULL]").await; + assert_eq!(row, "[{\"c0\":1},null]"); + + // Unsupported arrow field types surface as query errors instead of + // dropping the connection. + let error = client + .simple_query("SELECT struct(arrow_cast('1', 'Decimal256(38, 10)'))") + .await + .unwrap_err(); + let message = match error.as_db_error() { + Some(db_error) => db_error.message().to_string(), + None => error.to_string(), + }; + assert!( + message.contains("Unsupported arrow data type"), + "unexpected error message: {message}" + ); + + // The connection stays usable after the error. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_slow_query(store_type: StorageType) { let (mut guard, fe_pg_server) = setup_pg_server_with_slow_query_threshold( store_type, @@ -2441,3 +2672,184 @@ pub async fn test_declare_fetch_close_cursor(store_type: StorageType) { let _ = fe_pg_server.shutdown().await; guard.remove_all().await; } + +/// Keeps SQL and input rows fixed while changing only the selected batching protocols. +#[tokio::test(flavor = "multi_thread")] +async fn test_sql_batcher_alignment() { + use common_telemetry::{dump_metrics, info, init_default_ut_logging}; + use frontend::server::Services; + use frontend::service_config::BatcherOptions; + use servers::batcher::BatchingProtocol; + use servers::mysql::server::MYSQL_SERVER; + use tests_integration::standalone::GreptimeDbStandaloneBuilder; + + fn flushes() -> u64 { + dump_metrics() + .unwrap() + .lines() + .find_map(|line| { + line.strip_prefix("greptime_table_batcher_flush_total ") + .map(|value| value.parse().unwrap()) + }) + .unwrap_or(0) + } + + init_default_ut_logging(); + for protocols in [ + vec![], + vec![BatchingProtocol::Mysql], + vec![BatchingProtocol::Postgres], + vec![BatchingProtocol::Mysql, BatchingProtocol::Postgres], + ] { + let mysql_enabled = protocols.contains(&BatchingProtocol::Mysql); + let pg_enabled = protocols.contains(&BatchingProtocol::Postgres); + let mut instance = GreptimeDbStandaloneBuilder::new("sql_batcher_alignment") + .with_table_batcher(BatcherOptions { + protocols, + pending_rows_flush_interval: Duration::from_millis(10), + ..Default::default() + }) + .build() + .await; + let mut opts = instance.opts.clone(); + opts.http.addr = "127.0.0.1:0".into(); + opts.grpc.bind_addr = "127.0.0.1:0".into(); + opts.mysql.addr = "127.0.0.1:0".into(); + opts.postgres.addr = "127.0.0.1:0".into(); + let mut servers = Services::new(opts, instance.fe_instance().clone(), Default::default()) + .build() + .unwrap(); + servers.start_all().await.unwrap(); + let mysql_url = format!("mysql://{}/public", servers.addr(MYSQL_SERVER).unwrap()); + let mut mysql = MySqlConnection::connect(&mysql_url).await.unwrap(); + let (pg, connection) = tokio_postgres::connect( + &format!( + "postgres://{}/public", + servers.addr("POSTGRES_SERVER").unwrap() + ), + NoTls, + ) + .await + .unwrap(); + let pg_task = tokio::spawn(async move { connection.await.unwrap() }); + + for round in 0..2 { + mysql.execute("CREATE TABLE batch_sql (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val INT DEFAULT 7)").await.unwrap(); + let started = std::time::Instant::now(); + let before = flushes(); + // COM_QUERY and COM_STMT_EXECUTE must both opt in. + assert_eq!( + mysql + .execute("INSERT INTO batch_sql (ts, host) VALUES (0, 'a')") + .await + .unwrap() + .rows_affected(), + 1 + ); + assert_eq!( + sqlx::query("INSERT INTO batch_sql VALUES (0, ?, ?)") + .persistent(false) + .bind("b") + .bind(8_i32) + .execute(&mut mysql) + .await + .unwrap() + .rows_affected(), + 1 + ); + assert_eq!(flushes() - before, if mysql_enabled { 2 } else { 0 }); + + let before = flushes(); + let result = pg + .simple_query("INSERT INTO batch_sql (ts, host) VALUES (0, 'c')") + .await + .unwrap(); + assert!(matches!( + result.as_slice(), + [SimpleQueryMessage::CommandComplete(1)] + )); + let statement = pg + .prepare("INSERT INTO batch_sql VALUES (0, $1, $2)") + .await + .unwrap(); + assert_eq!(pg.execute(&statement, &[&"d", &9_i32]).await.unwrap(), 1); + assert_eq!(flushes() - before, if pg_enabled { 2 } else { 0 }); + + // Independent connections can submit to the same worker concurrently. + let (mysql_result, pg_result) = tokio::join!( + mysql.execute("INSERT INTO batch_sql VALUES (0, 'e', 10)"), + pg.simple_query("INSERT INTO batch_sql VALUES (0, 'f', 11)") + ); + assert_eq!(mysql_result.unwrap().rows_affected(), 1); + assert!(matches!( + pg_result.unwrap().as_slice(), + [SimpleQueryMessage::CommandComplete(1)] + )); + + // Errors must reach the client without affecting later requests. + assert!( + mysql + .execute("INSERT INTO batch_sql (missing) VALUES (1)") + .await + .is_err() + ); + assert!( + pg.simple_query("INSERT INTO batch_sql (missing) VALUES (1)") + .await + .is_err() + ); + let expected = vec![ + ("a".to_string(), 7_i32), + ("b".to_string(), 8), + ("c".to_string(), 7), + ("d".to_string(), 9), + ("e".to_string(), 10), + ("f".to_string(), 11), + ]; + // No polling: successful protocol responses guarantee visibility. + let rows: Vec<(String, i32)> = + sqlx::query_as("SELECT host, val FROM batch_sql ORDER BY host") + .persistent(false) + .fetch_all(&mut mysql) + .await + .unwrap(); + assert_eq!(rows, expected); + let rows: Vec<(String, i32)> = pg + .query("SELECT host, val FROM batch_sql ORDER BY host", &[]) + .await + .unwrap() + .iter() + .map(|row| (row.get(0), row.get(1))) + .collect(); + assert_eq!(rows, expected); + + mysql.execute("CREATE TABLE batch_copy (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val INT)").await.unwrap(); + assert_eq!( + mysql + .execute("INSERT INTO batch_copy SELECT * FROM batch_sql") + .await + .unwrap() + .rows_affected(), + 6 + ); + let rows: Vec<(String, i32)> = + sqlx::query_as("SELECT host, val FROM batch_copy ORDER BY host") + .persistent(false) + .fetch_all(&mut mysql) + .await + .unwrap(); + assert_eq!(rows, expected); + info!( + "sql_batcher_alignment mysql={mysql_enabled} postgres={pg_enabled} round={round} elapsed={:?}", + started.elapsed() + ); + mysql.execute("DROP TABLE batch_sql").await.unwrap(); + mysql.execute("DROP TABLE batch_copy").await.unwrap(); + } + mysql.close().await.unwrap(); + drop(pg); + pg_task.await.unwrap(); + servers.shutdown_all().await.unwrap(); + instance.guard.remove_all().await; + } +} diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 8bad66dae3c..6cce1736151 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -41,7 +41,9 @@ use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::{Channel, QueryContext}; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::test_util::setup_authenticated_grpc_database; +use tests_integration::test_util::{ + setup_authenticated_grpc_database, test_event_recorder_options, +}; use crate::event_recorder_test_util::{ assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, @@ -79,6 +81,7 @@ async fn test_table_ddl_procedure_events() { } let cluster = GreptimeDbClusterBuilder::new("table_ddl_procedure_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_metasrv_gc_config(gc_options) .with_datanode_gc_config(GcConfig { diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index b21838995c3..8d17e9b4e25 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -17,9 +17,9 @@ use std::sync::Arc; use common_test_util::temp_dir::create_temp_dir; use common_wal::config::DatanodeWalConfig; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ StorageType, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use uuid::Uuid; @@ -44,6 +44,7 @@ async fn test_view_ddl_events() { let home_dir = create_temp_dir("test_view_ddl_events_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_view_ddl_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) @@ -58,19 +59,6 @@ async fn test_view_ddl_events() { execute_view_ddl(instance, &source_table, &view).await; } -#[tokio::test(flavor = "multi_thread")] -async fn test_standalone_view_ddl_events() { - common_telemetry::init_default_ut_logging(); - let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_view_ddl_events") - .build() - .await; - let suffix = Uuid::new_v4().simple(); - let source_table = format!("view_ddl_event_source_{suffix}"); - let view = format!("view_ddl_event_{suffix}"); - - execute_view_ddl(standalone.fe_instance(), &source_table, &view).await; -} - async fn execute_view_ddl( instance: &Arc, source_table: &str, diff --git a/tests-integration/tests/wal_prune_event.rs b/tests-integration/tests/wal_prune_event.rs index db86ade4d9d..3879fb3e653 100644 --- a/tests-integration/tests/wal_prune_event.rs +++ b/tests-integration/tests/wal_prune_event.rs @@ -26,6 +26,7 @@ use meta_srv::procedure::wal_prune::{Context as WalPruneContext, WalPruneProcedu use rskafka::client::partition::{Compression, UnknownTopicHandling}; use rskafka::record::Record; use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::assert_single_event; @@ -35,6 +36,7 @@ async fn test_standalone_wal_prune_event() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_wal_prune_event") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let topic_name = format!("test_standalone_wal_prune_event-{}", uuid::Uuid::new_v4()); diff --git a/tests/cases/distributed/flow-tql/flow_tql.result b/tests/cases/distributed/flow-tql/flow_tql.result index 925fff3f242..73321d040e1 100644 --- a/tests/cases/distributed/flow-tql/flow_tql.result +++ b/tests/cases/distributed/flow-tql/flow_tql.result @@ -155,7 +155,7 @@ Error: 2000(InvalidSyntax), Invalid TQL syntax: sql parser error: Expected expre CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-(now()+'15s'::interval), '5s') count_values("status_code", http_requests); -Error: 3001(EngineExecuteQuery), Failed to convert float seconds to duration, raw: -15 +Error: 3001(EngineExecuteQuery), The end time must be greater than start time, start: 0, end: -15000 CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-now()+'15s'::interval, '5s') count_values("status_code", http_requests); @@ -228,14 +228,14 @@ SELECT * FROM cnt_reqs ORDER BY ts, status_code; +--------------------------+---------------------+-------------+ | count(http_requests.val) | ts | status_code | +--------------------------+---------------------+-------------+ -| 3.0 | 1970-01-01T00:00:00 | 200.0 | -| 1.0 | 1970-01-01T00:00:00 | 401.0 | -| 1.0 | 1970-01-01T00:00:05 | 401.0 | -| 2.0 | 1970-01-01T00:00:05 | 404.0 | -| 1.0 | 1970-01-01T00:00:05 | 500.0 | -| 2.0 | 1970-01-01T00:00:10 | 200.0 | -| 2.0 | 1970-01-01T00:00:10 | 201.0 | -| 4.0 | 1970-01-01T00:00:15 | 500.0 | +| 3.0 | 1970-01-01T00:00:00 | 200 | +| 1.0 | 1970-01-01T00:00:00 | 401 | +| 1.0 | 1970-01-01T00:00:05 | 401 | +| 2.0 | 1970-01-01T00:00:05 | 404 | +| 1.0 | 1970-01-01T00:00:05 | 500 | +| 2.0 | 1970-01-01T00:00:10 | 200 | +| 2.0 | 1970-01-01T00:00:10 | 201 | +| 4.0 | 1970-01-01T00:00:15 | 500 | +--------------------------+---------------------+-------------+ DROP FLOW calc_reqs; diff --git a/tests/cases/distributed/information_schema/ssts_limit.result b/tests/cases/distributed/information_schema/ssts_limit.result index 248593410fc..1b21a84c7af 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.result +++ b/tests/cases/distributed/information_schema/ssts_limit.result @@ -72,27 +72,28 @@ FROM ( LIMIT 1 ); -+---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | -| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | -| | Limit: skip=0, fetch=1 | -| | Projection: information_schema.ssts_manifest.region_id | -| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | -| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | -| | ]] | -| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | -| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | -| | CoalescePartitionsExec | -| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | -| | CoalescePartitionsExec: fetch=1 | -| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | -| | RepartitionExec: REDACTED -| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: [2], filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | -| | | -+---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | +| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | +| | Limit: skip=0, fetch=1 | +| | Projection: information_schema.ssts_manifest.region_id | +| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | +| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | +| | ]] | +| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | +| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec | +| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | +| | RepartitionExec: REDACTED +| | CoalescePartitionsExec: fetch=1 | +| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | +| | RepartitionExec: REDACTED +| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: [2], filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ DROP TABLE ssts_limit_case; diff --git a/tests/cases/standalone/common/aggregate/array_agg.result b/tests/cases/standalone/common/aggregate/array_agg.result index 5a73c289899..236a4582f14 100644 --- a/tests/cases/standalone/common/aggregate/array_agg.result +++ b/tests/cases/standalone/common/aggregate/array_agg.result @@ -144,3 +144,94 @@ DROP TABLE doubles; Affected Rows: 0 +-- On partitioned tables the aggregate is split into partial state and merge +CREATE TABLE array_agg_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO array_agg_partitioned VALUES + (1000, 1, 1), + (2000, 11, 2), + (3000, 21, 3), + (4000, 2, 4), + (5000, 12, 5), + (6000, 22, 6), + (7000, 3, 7); + +Affected Rows: 7 + +SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned; + ++-----------------------------------------------------------------------------------------+ +| array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] | ++-----------------------------------------------------------------------------------------+ +| [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] | ++-----------------------------------------------------------------------------------------+ + +SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g; + ++-------+-------------------------------------------------------------------------------------------+ +| g | array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] | ++-------+-------------------------------------------------------------------------------------------+ +| false | [3.0, 2.0, 1.0] | +| true | [7.0, 6.0, 5.0, 4.0] | ++-------+-------------------------------------------------------------------------------------------+ + +SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned; + ++---------------------------------------------------------------------------------------------------------------------------------------------+ +| array_agg(array_agg_partitioned.k) ORDER BY [array_agg_partitioned.k % Int64(10) ASC NULLS LAST, array_agg_partitioned.ts DESC NULLS FIRST] | ++---------------------------------------------------------------------------------------------------------------------------------------------+ +| [21, 11, 1, 22, 12, 2, 3] | ++---------------------------------------------------------------------------------------------------------------------------------------------+ + +-- nth_value needs sorted input, so it isn't split into partial state and merge +SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned; + ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ +| nth_value(array_agg_partitioned.lat,Int64(2)) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] | nth_value(array_agg_partitioned.lat,Int64(3)) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] | ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ +| 2.0 | 5.0 | ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ + +-- WITHIN GROUP aggregates don't need sorted input and are still split +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Aggregate: groupBy=[[]], aggr=[[__sum_merge(__sum_state(array_agg_partitioned.lat)) AS sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) AS approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Aggregate: groupBy=[[]], aggr=[[__sum_state(array_agg_partitioned.lat), __approx_percentile_cont_state(array_agg_partitioned.lat, Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]]]_| +|_|_TableScan: array_agg_partitioned_| +|_| ]]_| +| physical_plan | AggregateExec: mode=Final, gby=[], aggr=[__sum_merge(__sum_state(array_agg_partitioned.lat)) as sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) as approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]_| +|_|_CoalescePartitionsExec_| +|_|_AggregateExec: mode=Partial, gby=[], aggr=[__sum_merge(__sum_state(array_agg_partitioned.lat)) as sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) as approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]] | +|_|_RepartitionExec: partitioning=REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + ++--------------------------------+-------------------------------------------------------------------------------------------------+ +| sum(array_agg_partitioned.lat) | approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST] | ++--------------------------------+-------------------------------------------------------------------------------------------------+ +| 28.0 | 2.25 | ++--------------------------------+-------------------------------------------------------------------------------------------------+ + +DROP TABLE array_agg_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/array_agg.sql b/tests/cases/standalone/common/aggregate/array_agg.sql index dedabf1a18f..d0d7caa574f 100644 --- a/tests/cases/standalone/common/aggregate/array_agg.sql +++ b/tests/cases/standalone/common/aggregate/array_agg.sql @@ -54,3 +54,41 @@ DROP TABLE integers; DROP TABLE strings; DROP TABLE doubles; + +-- On partitioned tables the aggregate is split into partial state and merge +CREATE TABLE array_agg_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO array_agg_partitioned VALUES + (1000, 1, 1), + (2000, 11, 2), + (3000, 21, 3), + (4000, 2, 4), + (5000, 12, 5), + (6000, 22, 6), + (7000, 3, 7); + +SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned; + +SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g; + +SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned; + +-- nth_value needs sorted input, so it isn't split into partial state and merge +SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned; + +-- WITHIN GROUP aggregates don't need sorted input and are still split +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + +SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + +DROP TABLE array_agg_partitioned; diff --git a/tests/cases/standalone/common/aggregate/multi_regions.result b/tests/cases/standalone/common/aggregate/multi_regions.result index fa0bf929319..79ae85fbc04 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.result +++ b/tests/cases/standalone/common/aggregate/multi_regions.result @@ -113,6 +113,54 @@ select sum(val) from t group by idc; |_|_| Total rows: 0_| +-+-+-+ +insert into t values + (1000, 1, '1000-x', 'a'), + (1000, 2, '1000-y', 'a'), + (1000, 3, '2000-x', 'a'), + (1000, 4, '2000-y', 'a'); + +Affected Rows: 4 + +-- group keys derived from the partition column span regions +select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g; + ++---+----------+------------+ +| g | count(*) | sum(t.val) | ++---+----------+------------+ +| x | 2 | 4.0 | +| y | 2 | 6.0 | ++---+----------+------------+ + +-- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them +-- SQLNESS PROTOCOL POSTGRES +select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc; + ++--------+-----+------------+ +| host | idc | sum(t.val) | ++--------+-----+------------+ +| 1000-x | a | 1.0 | +| 1000-x | | 1.0 | +| 1000-y | a | 2.0 | +| 1000-y | | 2.0 | +| 2000-x | a | 3.0 | +| 2000-x | | 3.0 | +| 2000-y | a | 4.0 | +| 2000-y | | 4.0 | ++--------+-----+------------+ + +-- SQLNESS PROTOCOL POSTGRES +select host, sum(val) from t group by rollup(host) order by host; + ++--------+------------+ +| host | sum(t.val) | ++--------+------------+ +| 1000-x | 1.0 | +| 1000-y | 2.0 | +| 2000-x | 3.0 | +| 2000-y | 4.0 | +| | 10.0 | ++--------+------------+ + drop table t; Affected Rows: 0 diff --git a/tests/cases/standalone/common/aggregate/multi_regions.sql b/tests/cases/standalone/common/aggregate/multi_regions.sql index 3396bdddf19..4f99e1ed465 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.sql +++ b/tests/cases/standalone/common/aggregate/multi_regions.sql @@ -45,4 +45,20 @@ select sum(val) from t; explain analyze select sum(val) from t group by idc; +insert into t values + (1000, 1, '1000-x', 'a'), + (1000, 2, '1000-y', 'a'), + (1000, 3, '2000-x', 'a'), + (1000, 4, '2000-y', 'a'); + +-- group keys derived from the partition column span regions +select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g; + +-- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them +-- SQLNESS PROTOCOL POSTGRES +select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc; + +-- SQLNESS PROTOCOL POSTGRES +select host, sum(val) from t group by rollup(host) order by host; + drop table t; diff --git a/tests/cases/standalone/common/alter/alter_table_options.result b/tests/cases/standalone/common/alter/alter_table_options.result index 85e64db429f..09afb247055 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.result +++ b/tests/cases/standalone/common/alter/alter_table_options.result @@ -278,6 +278,36 @@ SHOW CREATE TABLE ato; | | ) | +-------+--------------------------------------------------------------+ +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +Affected Rows: 0 + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +Affected Rows: 0 + +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + -- SQLNESS ARG restart=true SHOW CREATE TABLE ato; @@ -303,6 +333,24 @@ SHOW CREATE TABLE ato; | | ) | +-------+--------------------------------------------------------------+ +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + ALTER TABLE ato UNSET 'compaction.twcs.active_window.trigger_file_num'; Affected Rows: 0 @@ -364,55 +412,6 @@ DROP TABLE ato; Affected Rows: 0 -CREATE TABLE write_buffer_size_options( - i INTEGER, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(i) -); - -Affected Rows: 0 - -ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; - -Affected Rows: 0 - -SHOW CREATE TABLE write_buffer_size_options; - -+---------------------------+----------------------------------------------------------+ -| Table | Create Table | -+---------------------------+----------------------------------------------------------+ -| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | -| | "i" INT NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("i") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | write_buffer_size = '128.0MiB' | -| | ) | -+---------------------------+----------------------------------------------------------+ - --- SQLNESS ARG restart=true -SHOW CREATE TABLE write_buffer_size_options; - -+---------------------------+----------------------------------------------------------+ -| Table | Create Table | -+---------------------------+----------------------------------------------------------+ -| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | -| | "i" INT NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("i") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | write_buffer_size = '128.0MiB' | -| | ) | -+---------------------------+----------------------------------------------------------+ - ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; Affected Rows: 0 diff --git a/tests/cases/standalone/common/alter/alter_table_options.sql b/tests/cases/standalone/common/alter/alter_table_options.sql index dd32ee0d4c3..de6470037b0 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.sql +++ b/tests/cases/standalone/common/alter/alter_table_options.sql @@ -62,9 +62,21 @@ ALTER TABLE ato SET 'compaction.twcs.inactive_window.l1_merge_trigger'='12'; SHOW CREATE TABLE ato; +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +SHOW CREATE TABLE write_buffer_size_options; + -- SQLNESS ARG restart=true SHOW CREATE TABLE ato; +SHOW CREATE TABLE write_buffer_size_options; + ALTER TABLE ato UNSET 'compaction.twcs.active_window.trigger_file_num'; ALTER TABLE ato UNSET 'compaction.twcs.active_window.l1_merge_trigger'; @@ -80,19 +92,6 @@ SHOW CREATE TABLE ato; DROP TABLE ato; -CREATE TABLE write_buffer_size_options( - i INTEGER, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(i) -); - -ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; - -SHOW CREATE TABLE write_buffer_size_options; - --- SQLNESS ARG restart=true -SHOW CREATE TABLE write_buffer_size_options; - ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; SHOW CREATE TABLE write_buffer_size_options; diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.result b/tests/cases/standalone/common/flow/flow_advance_ttl.result index 279a6552d21..27a28babebf 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.result +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.result @@ -116,7 +116,6 @@ SELECT number FROM distinct_basic; +--------+ +--------+ --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); +-------------------------------------+ @@ -356,7 +355,6 @@ SELECT number FROM distinct_basic; | 22 | +--------+ --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); +-------------------------------------+ @@ -416,3 +414,110 @@ DROP TABLE out_distinct_basic; Affected Rows: 0 +-- Streaming DISTINCT auto-sink retains all 1000 output tuples. +CREATE TABLE distinct_auto_sink ( + v INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX +) WITH ('ttl' = 'instant'); + +Affected Rows: 0 + +CREATE FLOW test_distinct_auto_sink SINK TO out_distinct_auto_sink AS +SELECT DISTINCT v FROM distinct_auto_sink; + +Affected Rows: 0 + +INSERT INTO distinct_auto_sink (v) +SELECT number FROM numbers LIMIT 1000; + +Affected Rows: 0 + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink'); + ++---------------------------------------------+ +| ADMIN FLUSH_FLOW('test_distinct_auto_sink') | ++---------------------------------------------+ +| FLOW_FLUSHED | ++---------------------------------------------+ + +SELECT count(*) AS rows, count(DISTINCT v) AS distinct_values, min(v) AS min_v, max(v) AS max_v +FROM out_distinct_auto_sink; + ++------+-----------------+-------+-------+ +| rows | distinct_values | min_v | max_v | ++------+-----------------+-------+-------+ +| 1000 | 1000 | 0 | 999 | ++------+-----------------+-------+-------+ + +DROP FLOW test_distinct_auto_sink; + +Affected Rows: 0 + +DROP TABLE distinct_auto_sink; + +Affected Rows: 0 + +DROP TABLE out_distinct_auto_sink; + +Affected Rows: 0 + +CREATE TABLE distinct_auto_sink_pk ( + v INT, + k INT, + extra_value INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX, + PRIMARY KEY (v, k) +) WITH ('ttl' = 'instant'); + +Affected Rows: 0 + +CREATE FLOW test_distinct_auto_sink_pk SINK TO out_distinct_auto_sink_pk AS +SELECT DISTINCT v, extra_value FROM distinct_auto_sink_pk; + +Affected Rows: 0 + +INSERT INTO distinct_auto_sink_pk (v, k, extra_value) +SELECT number % 10 AS v, number AS k, number AS extra_value FROM numbers LIMIT 1000; + +Affected Rows: 0 + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink_pk'); + ++------------------------------------------------+ +| ADMIN FLUSH_FLOW('test_distinct_auto_sink_pk') | ++------------------------------------------------+ +| FLOW_FLUSHED | ++------------------------------------------------+ + +SELECT + count(*) AS rows, + count(DISTINCT extra_value) AS distinct_values, + min(extra_value) AS min_value, + max(extra_value) AS max_value, + sum(CASE WHEN v = extra_value % 10 THEN 0 ELSE 1 END) AS mismatched_keys +FROM out_distinct_auto_sink_pk; + ++------+-----------------+-----------+-----------+-----------------+ +| rows | distinct_values | min_value | max_value | mismatched_keys | ++------+-----------------+-----------+-----------+-----------------+ +| 1000 | 1000 | 0 | 999 | 0 | ++------+-----------------+-----------+-----------+-----------------+ + +DROP FLOW test_distinct_auto_sink_pk; + +Affected Rows: 0 + +DROP TABLE distinct_auto_sink_pk; + +Affected Rows: 0 + +DROP TABLE out_distinct_auto_sink_pk; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.sql b/tests/cases/standalone/common/flow/flow_advance_ttl.sql index 653fa16637f..72356d96372 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.sql +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.sql @@ -51,7 +51,6 @@ FROM SELECT number FROM distinct_basic; --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); -- Recover the persisted streaming DISTINCT flow, then replan its first write @@ -158,7 +157,6 @@ FROM SELECT number FROM distinct_basic; --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); INSERT INTO @@ -181,3 +179,58 @@ SELECT number FROM distinct_basic; DROP FLOW test_distinct_basic; DROP TABLE distinct_basic; DROP TABLE out_distinct_basic; + +-- Streaming DISTINCT auto-sink retains all 1000 output tuples. +CREATE TABLE distinct_auto_sink ( + v INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX +) WITH ('ttl' = 'instant'); + +CREATE FLOW test_distinct_auto_sink SINK TO out_distinct_auto_sink AS +SELECT DISTINCT v FROM distinct_auto_sink; + +INSERT INTO distinct_auto_sink (v) +SELECT number FROM numbers LIMIT 1000; + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink'); + +SELECT count(*) AS rows, count(DISTINCT v) AS distinct_values, min(v) AS min_v, max(v) AS max_v +FROM out_distinct_auto_sink; + +DROP FLOW test_distinct_auto_sink; +DROP TABLE distinct_auto_sink; +DROP TABLE out_distinct_auto_sink; + +CREATE TABLE distinct_auto_sink_pk ( + v INT, + k INT, + extra_value INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX, + PRIMARY KEY (v, k) +) WITH ('ttl' = 'instant'); + +CREATE FLOW test_distinct_auto_sink_pk SINK TO out_distinct_auto_sink_pk AS +SELECT DISTINCT v, extra_value FROM distinct_auto_sink_pk; + +INSERT INTO distinct_auto_sink_pk (v, k, extra_value) +SELECT number % 10 AS v, number AS k, number AS extra_value FROM numbers LIMIT 1000; + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink_pk'); + +SELECT + count(*) AS rows, + count(DISTINCT extra_value) AS distinct_values, + min(extra_value) AS min_value, + max(extra_value) AS max_value, + sum(CASE WHEN v = extra_value % 10 THEN 0 ELSE 1 END) AS mismatched_keys +FROM out_distinct_auto_sink_pk; + +DROP FLOW test_distinct_auto_sink_pk; +DROP TABLE distinct_auto_sink_pk; +DROP TABLE out_distinct_auto_sink_pk; diff --git a/tests/cases/standalone/common/flow/flow_auto_sink_table.result b/tests/cases/standalone/common/flow/flow_auto_sink_table.result index 4ef3297125b..e2f28e0c441 100644 --- a/tests/cases/standalone/common/flow/flow_auto_sink_table.result +++ b/tests/cases/standalone/common/flow/flow_auto_sink_table.result @@ -114,16 +114,6 @@ GROUP BY Affected Rows: 0 --- SQLNESS ARG restart=true -SELECT 1; - -+----------+ -| Int64(1) | -+----------+ -| 1 | -+----------+ - --- SQLNESS SLEEP 3s SHOW CREATE TABLE out_num_cnt_basic; +-------------------+---------------------------------------------------+ diff --git a/tests/cases/standalone/common/flow/flow_auto_sink_table.sql b/tests/cases/standalone/common/flow/flow_auto_sink_table.sql index e67aa961910..87be556f66c 100644 --- a/tests/cases/standalone/common/flow/flow_auto_sink_table.sql +++ b/tests/cases/standalone/common/flow/flow_auto_sink_table.sql @@ -47,10 +47,7 @@ FROM numbers_input_basic GROUP BY ts; --- SQLNESS ARG restart=true -SELECT 1; --- SQLNESS SLEEP 3s SHOW CREATE TABLE out_num_cnt_basic; -- SQLNESS ARG restart=true diff --git a/tests/cases/standalone/common/flow/flow_basic.result b/tests/cases/standalone/common/flow/flow_basic.result index 5144c459ad6..f8f98e45dcd 100644 --- a/tests/cases/standalone/common/flow/flow_basic.result +++ b/tests/cases/standalone/common/flow/flow_basic.result @@ -447,190 +447,6 @@ DROP TABLE out_distinct_basic; Affected Rows: 0 -CREATE TABLE bytes_log ( - byte INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - -- event time - TIME INDEX(ts) -); - -Affected Rows: 0 - --- TODO(discord9): remove this after auto infer table's time index is impl -CREATE TABLE approx_rate ( - rate DOUBLE, - time_window TIMESTAMP, - update_at TIMESTAMP, - TIME INDEX(time_window) -); - -Affected Rows: 0 - -CREATE FLOW find_approx_rate SINK TO approx_rate AS -SELECT - (max(byte) - min(byte)) / 30.0 as rate, - date_bin(INTERVAL '30 second', ts) as time_window -from - bytes_log -GROUP BY - time_window; - -Affected Rows: 0 - -SHOW CREATE TABLE approx_rate; - -+-------------+--------------------------------------------+ -| Table | Create Table | -+-------------+--------------------------------------------+ -| approx_rate | CREATE TABLE IF NOT EXISTS "approx_rate" ( | -| | "rate" DOUBLE NULL, | -| | "time_window" TIMESTAMP(3) NOT NULL, | -| | "update_at" TIMESTAMP(3) NULL, | -| | TIME INDEX ("time_window") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+-------------+--------------------------------------------+ - -INSERT INTO - bytes_log -VALUES - (NULL, '2023-01-01 00:00:01'), - (300, '2023-01-01 00:00:29'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+------+---------------------+ -| rate | time_window | -+------+---------------------+ -| 0.0 | 2023-01-01T00:00:00 | -+------+---------------------+ - -INSERT INTO - bytes_log -VALUES - (NULL, '2022-01-01 00:00:01'), - (NULL, '2022-01-01 00:00:29'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+------+---------------------+ -| rate | time_window | -+------+---------------------+ -| | 2022-01-01T00:00:00 | -| 0.0 | 2023-01-01T00:00:00 | -+------+---------------------+ - -INSERT INTO - bytes_log -VALUES - (101, '2025-01-01 00:00:01'), - (300, '2025-01-01 00:00:29'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+-------------------+---------------------+ -| rate | time_window | -+-------------------+---------------------+ -| | 2022-01-01T00:00:00 | -| 0.0 | 2023-01-01T00:00:00 | -| 6.633333333333334 | 2025-01-01T00:00:00 | -+-------------------+---------------------+ - -INSERT INTO - bytes_log -VALUES - (450, '2025-01-01 00:00:32'), - (500, '2025-01-01 00:00:37'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+--------------------+---------------------+ -| rate | time_window | -+--------------------+---------------------+ -| | 2022-01-01T00:00:00 | -| 0.0 | 2023-01-01T00:00:00 | -| 6.633333333333334 | 2025-01-01T00:00:00 | -| 1.6666666666666667 | 2025-01-01T00:00:30 | -+--------------------+---------------------+ - -DROP TABLE bytes_log; - -Affected Rows: 0 - -DROP FLOW find_approx_rate; - -Affected Rows: 0 - -DROP TABLE approx_rate; - -Affected Rows: 0 - -- input table CREATE TABLE ngx_access_log ( client STRING, @@ -1774,6 +1590,14 @@ ADMIN FLUSH_FLOW('test_numbers_basic'); | FLOW_FLUSHED | +----------------------------------------+ +SELECT avg_after_filter_num FROM out_num_cnt_basic; + ++----------------------+ +| avg_after_filter_num | ++----------------------+ +| 0 | ++----------------------+ + DROP FLOW test_numbers_basic; Affected Rows: 0 diff --git a/tests/cases/standalone/common/flow/flow_basic.sql b/tests/cases/standalone/common/flow/flow_basic.sql index ef95ad88a9f..62e09db6d27 100644 --- a/tests/cases/standalone/common/flow/flow_basic.sql +++ b/tests/cases/standalone/common/flow/flow_basic.sql @@ -180,102 +180,6 @@ DROP TABLE distinct_basic; DROP TABLE out_distinct_basic; -CREATE TABLE bytes_log ( - byte INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - -- event time - TIME INDEX(ts) -); - --- TODO(discord9): remove this after auto infer table's time index is impl -CREATE TABLE approx_rate ( - rate DOUBLE, - time_window TIMESTAMP, - update_at TIMESTAMP, - TIME INDEX(time_window) -); - -CREATE FLOW find_approx_rate SINK TO approx_rate AS -SELECT - (max(byte) - min(byte)) / 30.0 as rate, - date_bin(INTERVAL '30 second', ts) as time_window -from - bytes_log -GROUP BY - time_window; - -SHOW CREATE TABLE approx_rate; - -INSERT INTO - bytes_log -VALUES - (NULL, '2023-01-01 00:00:01'), - (300, '2023-01-01 00:00:29'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -INSERT INTO - bytes_log -VALUES - (NULL, '2022-01-01 00:00:01'), - (NULL, '2022-01-01 00:00:29'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -INSERT INTO - bytes_log -VALUES - (101, '2025-01-01 00:00:01'), - (300, '2025-01-01 00:00:29'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -INSERT INTO - bytes_log -VALUES - (450, '2025-01-01 00:00:32'), - (500, '2025-01-01 00:00:37'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -DROP TABLE bytes_log; - -DROP FLOW find_approx_rate; - -DROP TABLE approx_rate; - -- input table CREATE TABLE ngx_access_log ( client STRING, @@ -827,6 +731,8 @@ VALUES -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_numbers_basic'); +SELECT avg_after_filter_num FROM out_num_cnt_basic; + DROP FLOW test_numbers_basic; DROP TABLE numbers_input_basic; DROP TABLE out_num_cnt_basic; diff --git a/tests/cases/standalone/common/flow/flow_call_df_func.result b/tests/cases/standalone/common/flow/flow_call_df_func.result index 960957a8e47..3bddade9825 100644 --- a/tests/cases/standalone/common/flow/flow_call_df_func.result +++ b/tests/cases/standalone/common/flow/flow_call_df_func.result @@ -15,15 +15,6 @@ SELECT sum(abs(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00' Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -50,15 +41,6 @@ SELECT "sum(abs(numbers_input_df_func.number))", time_window FROM out_num_cnt_df | 42 | 2021-07-01T00:00:00 | +----------------------------------------+---------------------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -114,15 +96,6 @@ SELECT abs(sum(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00' Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -148,15 +121,6 @@ SELECT "abs(sum(numbers_input_df_func.number))", time_window FROM out_num_cnt_df | 2 | 2021-07-01T00:00:00 | +----------------------------------------+---------------------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -211,15 +175,6 @@ SELECT max(number) - min(number) as maxmin, date_bin(INTERVAL '1 second', ts, '2 Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -244,15 +199,6 @@ SELECT maxmin, time_window FROM out_num_cnt_df_func; | 2 | 2021-07-01T00:00:00 | +--------+---------------------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -307,15 +253,6 @@ SELECT date_trunc('second', ts) as time_window, sum(number) as sum_num FROM numb Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -340,15 +277,6 @@ SELECT time_window, sum_num FROM out_num_cnt; | 2021-07-01T00:00:00 | 42 | +---------------------+---------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), diff --git a/tests/cases/standalone/common/flow/flow_call_df_func.sql b/tests/cases/standalone/common/flow/flow_call_df_func.sql index 21a09005960..d32f2b7b921 100644 --- a/tests/cases/standalone/common/flow/flow_call_df_func.sql +++ b/tests/cases/standalone/common/flow/flow_call_df_func.sql @@ -11,9 +11,6 @@ SINK TO out_num_cnt_df_func AS SELECT sum(abs(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00') as time_window FROM numbers_input_df_func GROUP BY time_window; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -26,9 +23,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); -- note that this quote-unquote column is a column-name, **not** a aggregation expr, generated by datafusion SELECT "sum(abs(numbers_input_df_func.number))", time_window FROM out_num_cnt_df_func; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -57,9 +51,6 @@ SINK TO out_num_cnt_df_func AS SELECT abs(sum(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00') as time_window FROM numbers_input_df_func GROUP BY time_window; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -71,9 +62,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); SELECT "abs(sum(numbers_input_df_func.number))", time_window FROM out_num_cnt_df_func; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -101,9 +89,6 @@ SINK TO out_num_cnt_df_func AS SELECT max(number) - min(number) as maxmin, date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00'::Timestamp) as time_window FROM numbers_input_df_func GROUP BY time_window; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -114,9 +99,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); SELECT maxmin, time_window FROM out_num_cnt_df_func; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -145,9 +127,6 @@ SINK TO out_num_cnt AS SELECT date_trunc('second', ts) as time_window, sum(number) as sum_num FROM numbers_input_df_func GROUP BY date_trunc('second', ts); --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -158,9 +137,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); SELECT time_window, sum_num FROM out_num_cnt; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), diff --git a/tests/cases/standalone/common/flow/flow_null.result b/tests/cases/standalone/common/flow/flow_null.result index 4ee94cfc099..a141b9e03ab 100644 --- a/tests/cases/standalone/common/flow/flow_null.result +++ b/tests/cases/standalone/common/flow/flow_null.result @@ -66,18 +66,6 @@ FROM | 600 | 2024-10-18T19:01:30 | +---------+---------------------+ --- Test if FLOWS table works, but don't care about the result since it vary from runs -SELECT - count(CASE WHEN state_size > 0 THEN 1 ELSE 0 END) as active_flows, -FROM - INFORMATION_SCHEMA.FLOWS; - -+--------------+ -| active_flows | -+--------------+ -| 1 | -+--------------+ - DROP FLOW requests_long_term; Affected Rows: 0 @@ -273,6 +261,8 @@ VALUES Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); @@ -302,6 +292,8 @@ VALUES Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); diff --git a/tests/cases/standalone/common/flow/flow_null.sql b/tests/cases/standalone/common/flow/flow_null.sql index 680d404b781..6c1219df2db 100644 --- a/tests/cases/standalone/common/flow/flow_null.sql +++ b/tests/cases/standalone/common/flow/flow_null.sql @@ -44,12 +44,6 @@ SELECT FROM sum_val_in_reqs; --- Test if FLOWS table works, but don't care about the result since it vary from runs -SELECT - count(CASE WHEN state_size > 0 THEN 1 ELSE 0 END) as active_flows, -FROM - INFORMATION_SCHEMA.FLOWS; - DROP FLOW requests_long_term; DROP TABLE sum_val_in_reqs; @@ -170,6 +164,8 @@ VALUES (2, "name2", "2024-10-18 19:00:00"), (3, "name3", "2024-10-18 19:00:00"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); @@ -183,6 +179,8 @@ VALUES ("name2", "2024-10-18 19:00:00"), ("name3", "2024-10-18 19:00:00"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); diff --git a/tests/cases/standalone/common/flow/flow_rebuild.result b/tests/cases/standalone/common/flow/flow_rebuild.result index 79c31fc8b11..a0b71b02c4f 100644 --- a/tests/cases/standalone/common/flow/flow_rebuild.result +++ b/tests/cases/standalone/common/flow/flow_rebuild.result @@ -349,73 +349,7 @@ SELECT wildcard FROM out_basic; | 2 | +----------+ -DROP TABLE input_basic; - -Affected Rows: 0 - -DROP TABLE out_basic; - -Affected Rows: 0 - -DROP FLOW test_wildcard_basic; - -Affected Rows: 0 - -- combination of different order of rebuild input table/flow -CREATE TABLE input_basic ( - "number" INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY(number), - TIME INDEX(ts) -); - -Affected Rows: 0 - -CREATE FLOW test_wildcard_basic sink TO out_basic EVAL INTERVAL '1m' AS -SELECT - COUNT(*) as wildcard -FROM - input_basic; - -Affected Rows: 0 - --- SQLNESS ARG restart=true -SELECT 1; - -+----------+ -| Int64(1) | -+----------+ -| 1 | -+----------+ - --- SQLNESS SLEEP 3s -INSERT INTO - input_basic -VALUES - (23, "2021-07-01 00:00:01.000"), - (24, "2021-07-01 00:00:01.500"); - -Affected Rows: 2 - --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_wildcard_basic'); - -+-----------------------------------------+ -| ADMIN FLUSH_FLOW('test_wildcard_basic') | -+-----------------------------------------+ -| FLOW_FLUSHED | -+-----------------------------------------+ - -SELECT wildcard FROM out_basic; - -+----------+ -| wildcard | -+----------+ -| 2 | -+----------+ - DROP TABLE input_basic; Affected Rows: 0 @@ -557,16 +491,6 @@ FROM Affected Rows: 0 --- SQLNESS ARG restart=true -SELECT 1; - -+----------+ -| Int64(1) | -+----------+ -| 1 | -+----------+ - --- SQLNESS SLEEP 3s INSERT INTO input_basic VALUES @@ -575,8 +499,6 @@ VALUES Affected Rows: 2 --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_wildcard_basic'); diff --git a/tests/cases/standalone/common/flow/flow_rebuild.sql b/tests/cases/standalone/common/flow/flow_rebuild.sql index d7d4f43749f..bdafed019d3 100644 --- a/tests/cases/standalone/common/flow/flow_rebuild.sql +++ b/tests/cases/standalone/common/flow/flow_rebuild.sql @@ -189,44 +189,8 @@ ADMIN FLUSH_FLOW('test_wildcard_basic'); SELECT wildcard FROM out_basic; -DROP TABLE input_basic; - -DROP TABLE out_basic; - -DROP FLOW test_wildcard_basic; - -- combination of different order of rebuild input table/flow -CREATE TABLE input_basic ( - "number" INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY(number), - TIME INDEX(ts) -); - -CREATE FLOW test_wildcard_basic sink TO out_basic EVAL INTERVAL '1m' AS -SELECT - COUNT(*) as wildcard -FROM - input_basic; - --- SQLNESS ARG restart=true -SELECT 1; - --- SQLNESS SLEEP 3s -INSERT INTO - input_basic -VALUES - (23, "2021-07-01 00:00:01.000"), - (24, "2021-07-01 00:00:01.500"); - --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_wildcard_basic'); - -SELECT wildcard FROM out_basic; - DROP TABLE input_basic; CREATE TABLE input_basic ( @@ -302,18 +266,12 @@ SELECT FROM input_basic; --- SQLNESS ARG restart=true -SELECT 1; - --- SQLNESS SLEEP 3s INSERT INTO input_basic VALUES (23, "2021-07-01 00:00:01.000"), (24, "2021-07-01 00:00:01.500"); --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_wildcard_basic'); diff --git a/tests/cases/standalone/common/flow/flow_view.result b/tests/cases/standalone/common/flow/flow_view.result index f823328d507..ab2577f93ae 100644 --- a/tests/cases/standalone/common/flow/flow_view.result +++ b/tests/cases/standalone/common/flow/flow_view.result @@ -10,12 +10,10 @@ INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512 Affected Rows: 4 --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); Affected Rows: 4 --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); Affected Rows: 4 diff --git a/tests/cases/standalone/common/flow/flow_view.sql b/tests/cases/standalone/common/flow/flow_view.sql index 420b9e84f4f..b10ef090908 100644 --- a/tests/cases/standalone/common/flow/flow_view.sql +++ b/tests/cases/standalone/common/flow/flow_view.sql @@ -4,10 +4,8 @@ CREATE TABLE user_agent_statistics ( user_agent STRING, total_count INT64, updat INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); CREATE FLOW user_agent_flow SINK TO user_agent_statistics EVAL INTERVAL '1m' AS SELECT user_agent, COUNT(user_agent) AS total_count FROM ngx_access_log GROUP BY user_agent; diff --git a/tests/cases/standalone/common/flow/show_create_flow.result b/tests/cases/standalone/common/flow/show_create_flow.result index f529cef1b9f..d451c552ec3 100644 --- a/tests/cases/standalone/common/flow/show_create_flow.result +++ b/tests/cases/standalone/common/flow/show_create_flow.result @@ -116,6 +116,8 @@ INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -155,6 +157,8 @@ INSERT INTO numbers_input_show VALUES (4,4),(5,4),(10, 3),(11, 4); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -195,6 +199,8 @@ INSERT INTO numbers_input_show VALUES (3, 1),(4, 2),(10, 3),(11, 4); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -324,6 +330,8 @@ INSERT INTO numbers_input_show VALUES(-4,0), (-3,1), (-2,2), (-1,3); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -384,6 +392,8 @@ INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -426,6 +436,8 @@ INSERT INTO numbers_input_show VALUES (10, 6),(11, 8),(15, 7),(18, 3); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); diff --git a/tests/cases/standalone/common/flow/show_create_flow.sql b/tests/cases/standalone/common/flow/show_create_flow.sql index 9856a6041f8..a2bcc4c06d9 100644 --- a/tests/cases/standalone/common/flow/show_create_flow.sql +++ b/tests/cases/standalone/common/flow/show_create_flow.sql @@ -45,6 +45,8 @@ SELECT flow_name, table_catalog, flow_definition, source_table_names FROM INFORM -- makesure it's not replaced in flownode INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -59,6 +61,8 @@ SELECT flow_name, table_catalog, flow_definition, source_table_names FROM INFORM -- makesure it's not replaced in flownode INSERT INTO numbers_input_show VALUES (4,4),(5,4),(10, 3),(11, 4); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -73,6 +77,8 @@ SELECT flow_name, table_catalog, flow_definition, source_table_names FROM INFORM -- makesure it's replaced in flownode INSERT INTO numbers_input_show VALUES (3, 1),(4, 2),(10, 3),(11, 4); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -118,6 +124,8 @@ SELECT * FROM out_num_cnt_show; INSERT INTO numbers_input_show VALUES(-4,0), (-3,1), (-2,2), (-1,3); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -145,6 +153,8 @@ CREATE FLOW filter_numbers_show SINK TO out_num_cnt_show AS SELECT number as n1 INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -160,6 +170,8 @@ SELECT flow_definition, source_table_names FROM INFORMATION_SCHEMA.FLOWS WHERE f INSERT INTO numbers_input_show VALUES (10, 6),(11, 8),(15, 7),(18, 3); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); diff --git a/tests/cases/standalone/common/function/admin/build_index_table.result b/tests/cases/standalone/common/function/admin/build_index_table.result deleted file mode 100644 index 4951df9e0da..00000000000 --- a/tests/cases/standalone/common/function/admin/build_index_table.result +++ /dev/null @@ -1,72 +0,0 @@ -CREATE TABLE test ( - ts TIMESTAMP TIME INDEX, - msg TEXT, -); - -Affected Rows: 0 - -INSERT INTO test VALUES -(1,"The quick brown fox jumps over the lazy dog"), -(2,"The quick brown fox jumps over the lazy cat"), -(3,"The quick brown fox jumps over the lazy mouse"), -(4,"The quick brown fox jumps over the lazy rabbit"), -(5,"The quick brown fox jumps over the lazy turtle"); - -Affected Rows: 5 - -SELECT * FROM test; - -+-------------------------+------------------------------------------------+ -| ts | msg | -+-------------------------+------------------------------------------------+ -| 1970-01-01T00:00:00.001 | The quick brown fox jumps over the lazy dog | -| 1970-01-01T00:00:00.002 | The quick brown fox jumps over the lazy cat | -| 1970-01-01T00:00:00.003 | The quick brown fox jumps over the lazy mouse | -| 1970-01-01T00:00:00.004 | The quick brown fox jumps over the lazy rabbit | -| 1970-01-01T00:00:00.005 | The quick brown fox jumps over the lazy turtle | -+-------------------------+------------------------------------------------+ - -ADMIN FLUSH_TABLE('test'); - -+---------------------------+ -| ADMIN FLUSH_TABLE('test') | -+---------------------------+ -| 0 | -+---------------------------+ - --- SQLNESS SLEEP 1s --- No fulltext index yet -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -+------------+ -| index_size | -+------------+ -| 0 | -+------------+ - -ALTER TABLE test MODIFY COLUMN msg SET FULLTEXT INDEX; - -Affected Rows: 0 - -ADMIN BUILD_INDEX('test'); - -+---------------------------+ -| ADMIN BUILD_INDEX('test') | -+---------------------------+ -| 0 | -+---------------------------+ - --- SQLNESS SLEEP 1s --- Fulltext index built -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -+------------+ -| index_size | -+------------+ -| 318 | -+------------+ - -DROP TABLE test; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/admin/build_index_table.sql b/tests/cases/standalone/common/function/admin/build_index_table.sql deleted file mode 100644 index da3b84f23da..00000000000 --- a/tests/cases/standalone/common/function/admin/build_index_table.sql +++ /dev/null @@ -1,29 +0,0 @@ -CREATE TABLE test ( - ts TIMESTAMP TIME INDEX, - msg TEXT, -); - -INSERT INTO test VALUES -(1,"The quick brown fox jumps over the lazy dog"), -(2,"The quick brown fox jumps over the lazy cat"), -(3,"The quick brown fox jumps over the lazy mouse"), -(4,"The quick brown fox jumps over the lazy rabbit"), -(5,"The quick brown fox jumps over the lazy turtle"); - -SELECT * FROM test; - -ADMIN FLUSH_TABLE('test'); - --- SQLNESS SLEEP 1s --- No fulltext index yet -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -ALTER TABLE test MODIFY COLUMN msg SET FULLTEXT INDEX; - -ADMIN BUILD_INDEX('test'); - --- SQLNESS SLEEP 1s --- Fulltext index built -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -DROP TABLE test; diff --git a/tests/cases/standalone/common/function/admin/build_index_table_restart.result b/tests/cases/standalone/common/function/admin/build_index_table_restart.result index 78451e3121c..f6e02582189 100644 --- a/tests/cases/standalone/common/function/admin/build_index_table_restart.result +++ b/tests/cases/standalone/common/function/admin/build_index_table_restart.result @@ -14,6 +14,18 @@ INSERT INTO build_index_restart_test VALUES Affected Rows: 5 +SELECT * FROM build_index_restart_test; + ++-------------------------+------------------------------------------------+ +| ts | msg | ++-------------------------+------------------------------------------------+ +| 1970-01-01T00:00:00.001 | The quick brown fox jumps over the lazy dog | +| 1970-01-01T00:00:00.002 | The quick brown fox jumps over the lazy cat | +| 1970-01-01T00:00:00.003 | The quick brown fox jumps over the lazy mouse | +| 1970-01-01T00:00:00.004 | The quick brown fox jumps over the lazy rabbit | +| 1970-01-01T00:00:00.005 | The quick brown fox jumps over the lazy turtle | ++-------------------------+------------------------------------------------+ + ADMIN FLUSH_TABLE('build_index_restart_test'); +-----------------------------------------------+ @@ -32,6 +44,21 @@ SHOW INDEX FROM build_index_restart_test; | build_index_restart_test | 1 | TIME INDEX | 1 | ts | A | | | | | TIME | | | YES | | +--------------------------+------------+------------+--------------+-------------+-----------+-------------+----------+--------+------+------------+---------+---------------+---------+------------+ +-- Index size is 0 before the fulltext index is built +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + ++------------+ +| index_size | ++------------+ +| 0 | ++------------+ + -- No physical fulltext index metadata yet. -- ssts_index_meta.index_type uses physical backend names (fulltext_bloom/fulltext_tantivy), -- unlike SHOW INDEX's greptime-fulltext-index-* display value. @@ -108,6 +135,21 @@ AND index_type LIKE 'fulltext%'; | 1 | +---------------------------+ +-- Fulltext index size is reported after the build +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + ++------------+ +| index_size | ++------------+ +| 318 | ++------------+ + -- SQLNESS ARG restart=true -- After restart, verify fulltext query still works SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; diff --git a/tests/cases/standalone/common/function/admin/build_index_table_restart.sql b/tests/cases/standalone/common/function/admin/build_index_table_restart.sql index b5503ed18e2..745e98eee2c 100644 --- a/tests/cases/standalone/common/function/admin/build_index_table_restart.sql +++ b/tests/cases/standalone/common/function/admin/build_index_table_restart.sql @@ -10,12 +10,23 @@ INSERT INTO build_index_restart_test VALUES (4,"The quick brown fox jumps over the lazy rabbit"), (5,"The quick brown fox jumps over the lazy turtle"); +SELECT * FROM build_index_restart_test; + ADMIN FLUSH_TABLE('build_index_restart_test'); -- SQLNESS SLEEP 1s -- No fulltext index yet SHOW INDEX FROM build_index_restart_test; +-- Index size is 0 before the fulltext index is built +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + -- No physical fulltext index metadata yet. -- ssts_index_meta.index_type uses physical backend names (fulltext_bloom/fulltext_tantivy), -- unlike SHOW INDEX's greptime-fulltext-index-* display value. @@ -55,6 +66,15 @@ WHERE table_id = ( ) AND index_type LIKE 'fulltext%'; +-- Fulltext index size is reported after the build +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + -- SQLNESS ARG restart=true -- After restart, verify fulltext query still works SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; diff --git a/tests/cases/standalone/common/function/ai_functions.result b/tests/cases/standalone/common/function/ai_functions.result new file mode 100644 index 00000000000..057c9ecd990 --- /dev/null +++ b/tests/cases/standalone/common/function/ai_functions.result @@ -0,0 +1,144 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT ai_match(NULL, 'The event reports a failed payment.') AS score; + ++-------+ +| score | ++-------+ +| | ++-------+ + +SELECT ai_match('message', NULL) AS score; + ++-------+ +| score | ++-------+ +| | ++-------+ + +SELECT arrow_typeof(ai_match(NULL, 'condition')) AS score_type; + ++------------+ +| score_type | ++------------+ +| Float64 | ++------------+ + +SELECT ai_choose(NULL, 'Route the ticket', '{"billing":null}') AS null_text, + ai_choose('message', NULL, 'invalid JSON') AS null_prompt, + ai_choose('message', 'Route the ticket', NULL) AS null_criteria; + ++-----------+-------------+---------------+ +| null_text | null_prompt | null_criteria | ++-----------+-------------+---------------+ +| | | | ++-----------+-------------+---------------+ + +SELECT ai_score(NULL, 'Rate severity', '["low","high"]') AS null_text, + ai_score('message', NULL, 'invalid JSON') AS null_prompt, + ai_score('message', 'Rate severity', NULL) AS null_criteria; + ++-----------+-------------+---------------+ +| null_text | null_prompt | null_criteria | ++-----------+-------------+---------------+ +| | | | ++-----------+-------------+---------------+ + +SELECT arrow_typeof(ai_choose(NULL, 'prompt', '{"billing":null}')) AS choice_type, + arrow_typeof(ai_score(NULL, 'prompt', '["low","high"]')) AS score_type; + ++-------------+------------+ +| choice_type | score_type | ++-------------+------------+ +| Utf8 | BinaryView | ++-------------+------------+ + +SELECT json_to_string(rating) AS rating, + json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[2]') AS high_probability +FROM (SELECT ai_score(NULL, 'prompt', '["low","medium","high"]') AS rating) AS rated; + ++--------+-------+------------+------------------+ +| rating | score | confidence | high_probability | ++--------+-------+------------+------------------+ +| | | | | ++--------+-------+------------+------------------+ + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE ai_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +Affected Rows: 0 + +INSERT INTO ai_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +Affected Rows: 2 + +SELECT occurred_at, service, message FROM ai_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match((message), 'The event reports that a payment still failed after retries.') >= 0.8 +ORDER BY occurred_at; + ++-------------+---------+---------+ +| occurred_at | service | message | ++-------------+---------+---------+ ++-------------+---------+---------+ + +SELECT occurred_at, + ai_choose(message, 'Route the ticket', '{"billing":"Payments","technical":"Errors"}') AS team, + ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating +FROM ai_events +ORDER BY occurred_at; + ++---------------------+------+--------+ +| occurred_at | team | rating | ++---------------------+------+--------+ +| 2026-09-19T01:00:00 | | | +| 2026-09-19T02:00:00 | | | ++---------------------+------+--------+ + +DROP TABLE ai_events; + +Affected Rows: 0 + +-- Matching uses two arguments; thresholds are expressed as SQL comparisons. +SELECT ai_match('message', 'condition', 0.8); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_match' function: coercion from Utf8, Utf8, Float64 to the signature Exact(Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_match(Utf8, Utf8, Float64)'. You might need to add explicit type casts. + Candidate functions: + ai_match(Utf8View, Utf8View) + +-- Invalid criteria fail locally before any HTTP request. +SELECT ai_choose('message', 'prompt', 'not json'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_choose criteria is not valid JSON: expected ident at line 1 column 2 + +SELECT ai_choose('message', 'prompt', '{}'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_choose criteria must be a JSON object with 1 to 255 options; descriptions must be strings, objects, arrays, or null + +SELECT ai_score('message', 'prompt', '["only one level"]'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_score criteria must be a JSON array with 2 to 10 levels; descriptions must be strings, objects, or arrays + +-- Both additional modes require criteria. +SELECT ai_choose('message', 'prompt'); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_choose' function: coercion from Utf8, Utf8 to the signature Exact(Utf8View, Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_choose(Utf8, Utf8)'. You might need to add explicit type casts. + Candidate functions: + ai_choose(Utf8View, Utf8View, Utf8View) + +SELECT ai_score('message', 'prompt'); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_score' function: coercion from Utf8, Utf8 to the signature Exact(Utf8View, Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_score(Utf8, Utf8)'. You might need to add explicit type casts. + Candidate functions: + ai_score(Utf8View, Utf8View, Utf8View) + diff --git a/tests/cases/standalone/common/function/ai_functions.sql b/tests/cases/standalone/common/function/ai_functions.sql new file mode 100644 index 00000000000..fd3e80d1e18 --- /dev/null +++ b/tests/cases/standalone/common/function/ai_functions.sql @@ -0,0 +1,65 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT ai_match(NULL, 'The event reports a failed payment.') AS score; + +SELECT ai_match('message', NULL) AS score; + +SELECT arrow_typeof(ai_match(NULL, 'condition')) AS score_type; + +SELECT ai_choose(NULL, 'Route the ticket', '{"billing":null}') AS null_text, + ai_choose('message', NULL, 'invalid JSON') AS null_prompt, + ai_choose('message', 'Route the ticket', NULL) AS null_criteria; + +SELECT ai_score(NULL, 'Rate severity', '["low","high"]') AS null_text, + ai_score('message', NULL, 'invalid JSON') AS null_prompt, + ai_score('message', 'Rate severity', NULL) AS null_criteria; + +SELECT arrow_typeof(ai_choose(NULL, 'prompt', '{"billing":null}')) AS choice_type, + arrow_typeof(ai_score(NULL, 'prompt', '["low","high"]')) AS score_type; + +SELECT json_to_string(rating) AS rating, + json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[2]') AS high_probability +FROM (SELECT ai_score(NULL, 'prompt', '["low","medium","high"]') AS rating) AS rated; + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE ai_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +INSERT INTO ai_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +SELECT occurred_at, service, message FROM ai_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match((message), 'The event reports that a payment still failed after retries.') >= 0.8 +ORDER BY occurred_at; + +SELECT occurred_at, + ai_choose(message, 'Route the ticket', '{"billing":"Payments","technical":"Errors"}') AS team, + ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating +FROM ai_events +ORDER BY occurred_at; + +DROP TABLE ai_events; + +-- Matching uses two arguments; thresholds are expressed as SQL comparisons. +SELECT ai_match('message', 'condition', 0.8); + +-- Invalid criteria fail locally before any HTTP request. +SELECT ai_choose('message', 'prompt', 'not json'); + +SELECT ai_choose('message', 'prompt', '{}'); + +SELECT ai_score('message', 'prompt', '["only one level"]'); + +-- Both additional modes require criteria. +SELECT ai_choose('message', 'prompt'); + +SELECT ai_score('message', 'prompt'); diff --git a/tests/cases/standalone/common/function/geo.result b/tests/cases/standalone/common/function/geo.result index 31efbf401da..3fc73a3b98b 100644 --- a/tests/cases/standalone/common/function/geo.result +++ b/tests/cases/standalone/common/function/geo.result @@ -437,3 +437,47 @@ FROM | true | false | true | false | false | true | +--------------------------+--------------------------+------------------------+------------------------+----------------------------------+----------------------------------+ +CREATE TABLE geo_path_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + lon DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO geo_path_partitioned VALUES + (1000, 1, 1, 11), + (2000, 11, 2, 12), + (3000, 21, 3, 13), + (4000, 2, 4, 14), + (5000, 12, 5, 15), + (6000, 22, 6, 16), + (7000, 3, 7, 17); + +Affected Rows: 7 + +SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + ++-------+-------------------------------------------------------------------------------------+ +| g | geo_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) | ++-------+-------------------------------------------------------------------------------------+ +| false | {lat: [1.0, 2.0, 3.0], lng: [11.0, 12.0, 13.0]} | +| true | {lat: [4.0, 5.0, 6.0, 7.0], lng: [14.0, 15.0, 16.0, 17.0]} | ++-------+-------------------------------------------------------------------------------------+ + +SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + ++-------+---------------------------------------------------------------------------------------------+ +| g | json_encode_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) | ++-------+---------------------------------------------------------------------------------------------+ +| false | [[11.0,1.0],[12.0,2.0],[13.0,3.0]] | +| true | [[14.0,4.0],[15.0,5.0],[16.0,6.0],[17.0,7.0]] | ++-------+---------------------------------------------------------------------------------------------+ + +DROP TABLE geo_path_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/geo.sql b/tests/cases/standalone/common/function/geo.sql index 43025cc196b..1de7d7c82aa 100644 --- a/tests/cases/standalone/common/function/geo.sql +++ b/tests/cases/standalone/common/function/geo.sql @@ -180,3 +180,27 @@ FROM 'POLYGON ((-121.491698 38.653343, -121.582353 38.556757, -121.469721 38.449287, -121.315883 38.541721, -121.491698 38.653343))' AS polygon2, 'POLYGON ((-122.089628 37.450332, -122.20535 37.378342, -122.093062 37.36088, -122.044301 37.372886, -122.089628 37.450332))' AS polygon3, ); + +CREATE TABLE geo_path_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + lon DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO geo_path_partitioned VALUES + (1000, 1, 1, 11), + (2000, 11, 2, 12), + (3000, 21, 3, 13), + (4000, 2, 4, 14), + (5000, 12, 5, 15), + (6000, 22, 6, 16), + (7000, 3, 7, 17); + +SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + +SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + +DROP TABLE geo_path_partitioned; diff --git a/tests/cases/standalone/common/function/jev.result b/tests/cases/standalone/common/function/jev.result deleted file mode 100644 index 4f6d91d2867..00000000000 --- a/tests/cases/standalone/common/function/jev.result +++ /dev/null @@ -1,66 +0,0 @@ --- No credentials or external service needed: NULL propagates without an API call. -SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - -SELECT jev('message', NULL, 0.8) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - -SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - --- Validate SQL registration, coercion, and asynchronous filtering on a real table. -CREATE TABLE jev_events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); - -Affected Rows: 0 - -INSERT INTO jev_events VALUES - ('2026-09-19T01:00:00Z', 'payments', NULL), - ('2026-09-19T02:00:00Z', 'auth', NULL); - -Affected Rows: 2 - -SELECT occurred_at, service, message FROM jev_events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) -ORDER BY occurred_at; - -+-------------+---------+---------+ -| occurred_at | service | message | -+-------------+---------+---------+ -+-------------+---------+---------+ - -DROP TABLE jev_events; - -Affected Rows: 0 - --- Invalid thresholds fail locally before any HTTP request. -SELECT jev('message', 'condition', -0.1); - -Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] - -SELECT jev('message', 'condition', 1.1); - -Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] - diff --git a/tests/cases/standalone/common/function/jev.sql b/tests/cases/standalone/common/function/jev.sql deleted file mode 100644 index f4c252c9273..00000000000 --- a/tests/cases/standalone/common/function/jev.sql +++ /dev/null @@ -1,32 +0,0 @@ --- No credentials or external service needed: NULL propagates without an API call. -SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; - -SELECT jev('message', NULL, 0.8) AS matched; - -SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; - --- Validate SQL registration, coercion, and asynchronous filtering on a real table. -CREATE TABLE jev_events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); - -INSERT INTO jev_events VALUES - ('2026-09-19T01:00:00Z', 'payments', NULL), - ('2026-09-19T02:00:00Z', 'auth', NULL); - -SELECT occurred_at, service, message FROM jev_events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) -ORDER BY occurred_at; - -DROP TABLE jev_events; - --- Invalid thresholds fail locally before any HTTP request. -SELECT jev('message', 'condition', -0.1); - -SELECT jev('message', 'condition', 1.1); diff --git a/tests/cases/standalone/common/function/vector/vector.result b/tests/cases/standalone/common/function/vector/vector.result index c546f9ad25e..06f7c684fa1 100644 --- a/tests/cases/standalone/common/function/vector/vector.result +++ b/tests/cases/standalone/common/function/vector/vector.result @@ -448,3 +448,89 @@ FROM ( | [4,5,6,10,-8] | +---------------------------------------------------+ +SELECT h, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM ( + SELECT 'a' AS h, '[1.0, 1.0]' AS v + UNION ALL + SELECT 'a' AS h, '[2.0, 2.0]' AS v + UNION ALL + SELECT 'b' AS h, '[3.0, 3.0]' AS v +) GROUP BY h ORDER BY h; + ++---+---------------------------+---------------------------+-------------------------------+ +| h | vec_to_string(vec_sum(v)) | vec_to_string(vec_avg(v)) | vec_to_string(vec_product(v)) | ++---+---------------------------+---------------------------+-------------------------------+ +| a | [3,3] | [1.5,1.5] | [2,2] | +| b | [3,3] | [3,3] | [3,3] | ++---+---------------------------+---------------------------+-------------------------------+ + +-- On partitioned tables the aggregates are split into partial state and merge. +-- Rows only land in two of the three partitions, the third one has an empty state. +-- The two non-empty partitions differ in row count and mean, so vec_avg must weight by count. +CREATE TABLE vector_aggr_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + g STRING, + v VECTOR(2), + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO vector_aggr_partitioned VALUES + (1000, 1, 'a', '[1.0, 1.0]'), + (2000, 11, 'a', '[8.0, 8.0]'), + (3000, 2, 'a', '[3.0, 3.0]'), + (3500, 3, 'a', '[5.0, 5.0]'), + (4000, 12, 'b', '[4.0, 4.0]'); + +Affected Rows: 5 + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| [21,21] | [4.2,4.2] | [480,480] | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| g | vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| a | [17,17] | [4.25,4.25] | [120,120] | +| b | [4,4] | [4,4] | [4,4] | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +-- A NULL vector makes vec_sum and vec_product NULL, vec_avg skips it. +INSERT INTO vector_aggr_partitioned VALUES (5000, 21, 'b', NULL); + +Affected Rows: 1 + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| | [4.2,4.2] | | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| g | vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| a | [17,17] | [4.25,4.25] | [120,120] | +| b | | [4,4] | | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +DROP TABLE vector_aggr_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/vector/vector.sql b/tests/cases/standalone/common/function/vector/vector.sql index 9bbf1583f53..c53438c854d 100644 --- a/tests/cases/standalone/common/function/vector/vector.sql +++ b/tests/cases/standalone/common/function/vector/vector.sql @@ -150,3 +150,48 @@ FROM ( UNION ALL SELECT '[4.0, 5.0, 6.0, 10, -8, 100]' AS v ) ORDER BY v; + +SELECT h, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM ( + SELECT 'a' AS h, '[1.0, 1.0]' AS v + UNION ALL + SELECT 'a' AS h, '[2.0, 2.0]' AS v + UNION ALL + SELECT 'b' AS h, '[3.0, 3.0]' AS v +) GROUP BY h ORDER BY h; + +-- On partitioned tables the aggregates are split into partial state and merge. +-- Rows only land in two of the three partitions, the third one has an empty state. +-- The two non-empty partitions differ in row count and mean, so vec_avg must weight by count. +CREATE TABLE vector_aggr_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + g STRING, + v VECTOR(2), + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO vector_aggr_partitioned VALUES + (1000, 1, 'a', '[1.0, 1.0]'), + (2000, 11, 'a', '[8.0, 8.0]'), + (3000, 2, 'a', '[3.0, 3.0]'), + (3500, 3, 'a', '[5.0, 5.0]'), + (4000, 12, 'b', '[4.0, 4.0]'); + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + +-- A NULL vector makes vec_sum and vec_product NULL, vec_avg skips it. +INSERT INTO vector_aggr_partitioned VALUES (5000, 21, 'b', NULL); + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + +DROP TABLE vector_aggr_partitioned; diff --git a/tests/cases/standalone/common/function/vector/vector_index.result b/tests/cases/standalone/common/function/vector/vector_index.result deleted file mode 100644 index 44c6e447054..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index.result +++ /dev/null @@ -1,602 +0,0 @@ --- Test vector index creation and KNN search --- ============================================ --- Part 1: Basic L2 squared distance tests --- ============================================ --- Create a table with vector column and L2sq vector index -CREATE TABLE vectors_l2sq ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert test vectors -INSERT INTO vectors_l2sq VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.0, 0.0, 1.0]'), - (5, '2024-01-01 00:00:04', '[1.0, 1.0, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.0, 1.0, 1.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.0, 0.0, 1.0, 1.0]'), - (8, '2024-01-01 00:00:07', '[1.0, 0.0, 0.0, 1.0]'); - -Affected Rows: 8 - --- Query BEFORE flush (memtable search) --- Expected: vec_id=1 (distance=0), vec_id=5 (distance=1), vec_id=8 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 5 | 1.0 | -| 8 | 1.0 | -+--------+----------+ - --- Flush to create SST files with vector index -ADMIN FLUSH_TABLE('vectors_l2sq'); - -+-----------------------------------+ -| ADMIN FLUSH_TABLE('vectors_l2sq') | -+-----------------------------------+ -| 0 | -+-----------------------------------+ - --- Query AFTER flush (SST index search) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 5 | 1.0 | -| 8 | 1.0 | -+--------+----------+ - --- Query with different target vector --- Expected: vec_id=6 (distance=0), vec_id=2 (distance=1), vec_id=3 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 6 | 0.0 | -| 2 | 1.0 | -| 3 | 1.0 | -+--------+----------+ - -DROP TABLE vectors_l2sq; - -Affected Rows: 0 - --- ============================================ --- Part 2: Cosine distance tests --- ============================================ -CREATE TABLE vectors_cosine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert vectors with different magnitudes but same/different directions -INSERT INTO vectors_cosine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search with cosine distance --- vec_id=1 and vec_id=2 should have same cosine distance (0) since they point same direction -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+------------+ -| vec_id | distance | -+--------+------------+ -| 1 | 0.0 | -| 2 | 0.0 | -| 4 | 0.29289323 | -+--------+------------+ - -ADMIN FLUSH_TABLE('vectors_cosine'); - -+-------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_cosine') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SST index search with cosine distance -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+------------+ -| vec_id | distance | -+--------+------------+ -| 1 | 0.0 | -| 2 | 0.0 | -| 4 | 0.29289323 | -+--------+------------+ - -DROP TABLE vectors_cosine; - -Affected Rows: 0 - --- ============================================ --- Part 3: Dot product (inner product) tests --- ============================================ -CREATE TABLE vectors_dot ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'dot'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_dot VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search with dot product --- Larger dot product means more similar, so we use negative for ordering --- vec_id=2 should be best (dot=2), vec_id=1 and vec_id=4 have dot=1 -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -+--------+-------------+ -| vec_id | dot_product | -+--------+-------------+ -| 2 | 2.0 | -| 1 | 1.0 | -| 4 | 1.0 | -+--------+-------------+ - -ADMIN FLUSH_TABLE('vectors_dot'); - -+----------------------------------+ -| ADMIN FLUSH_TABLE('vectors_dot') | -+----------------------------------+ -| 0 | -+----------------------------------+ - --- SST index search with dot product -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -+--------+-------------+ -| vec_id | dot_product | -+--------+-------------+ -| 2 | 2.0 | -| 1 | 1.0 | -| 4 | 1.0 | -+--------+-------------+ - -DROP TABLE vectors_dot; - -Affected Rows: 0 - --- ============================================ --- Part 4: NULL vector handling tests --- ============================================ -CREATE TABLE vectors_null ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert vectors with some NULLs -INSERT INTO vectors_null VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', NULL), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', NULL), - (5, '2024-01-01 00:00:04', '[0.0, 0.0, 1.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search should skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 3 | 2.0 | -| 5 | 2.0 | -+--------+----------+ - -ADMIN FLUSH_TABLE('vectors_null'); - -+-----------------------------------+ -| ADMIN FLUSH_TABLE('vectors_null') | -+-----------------------------------+ -| 0 | -+-----------------------------------+ - --- SST index search should also skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 3 | 2.0 | -| 5 | 2.0 | -+--------+----------+ - -DROP TABLE vectors_null; - -Affected Rows: 0 - --- ============================================ --- Part 5: Mixed memtable + SST search tests --- ============================================ -CREATE TABLE vectors_mixed ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert first batch and flush to SST -INSERT INTO vectors_mixed VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'); - -Affected Rows: 3 - -ADMIN FLUSH_TABLE('vectors_mixed'); - -+------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_mixed') | -+------------------------------------+ -| 0 | -+------------------------------------+ - --- Insert second batch (stays in memtable) -INSERT INTO vectors_mixed VALUES - (4, '2024-01-01 00:00:03', '[0.5, 0.5, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.9, 0.1, 0.0, 0.0]'); - -Affected Rows: 2 - --- Query should search both SST (vec_id 1,2,3) and memtable (vec_id 4,5) --- Target: [1.0, 0.0, 0.0, 0.0] --- Expected: vec_id=1 (dist=0), vec_id=5 (dist=0.02), vec_id=4 (dist=0.5) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_mixed -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+-------------+ -| vec_id | distance | -+--------+-------------+ -| 1 | 0.0 | -| 5 | 0.020000005 | -| 4 | 0.5 | -+--------+-------------+ - -DROP TABLE vectors_mixed; - -Affected Rows: 0 - --- ============================================ --- Part 6: KNN with WHERE clause tests --- ============================================ -CREATE TABLE vectors_filter ( - vec_id INT, - "category" STRING, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_filter VALUES - (1, 'A', '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, 'B', '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, 'A', '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, 'B', '2024-01-01 00:00:03', '[0.1, 0.9, 0.0, 0.0]'), - (5, 'A', '2024-01-01 00:00:04', '[0.5, 0.5, 0.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+----------+----------+ -| vec_id | category | distance | -+--------+----------+----------+ -| 1 | A | 0.0 | -| 5 | A | 0.5 | -+--------+----------+----------+ - -ADMIN FLUSH_TABLE('vectors_filter'); - -+-------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_filter') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SST index search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+----------+----------+ -| vec_id | category | distance | -+--------+----------+----------+ -| 1 | A | 0.0 | -| 5 | A | 0.5 | -+--------+----------+----------+ - --- Filter with time range -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE ts >= '2024-01-01 00:00:02' AND ts <= '2024-01-01 00:00:04' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+-----------+ -| vec_id | distance | -+--------+-----------+ -| 5 | 0.5 | -| 4 | 1.6199999 | -+--------+-----------+ - -DROP TABLE vectors_filter; - -Affected Rows: 0 - --- ============================================ --- Part 7: Higher dimension vectors --- ============================================ -CREATE TABLE vectors_high_dim ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(128) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert high-dimensional vectors (simplified: first few elements differ) -INSERT INTO vectors_high_dim VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'); - -Affected Rows: 3 - --- Memtable search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -ADMIN FLUSH_TABLE('vectors_high_dim'); - -+---------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_high_dim') | -+---------------------------------------+ -| 0 | -+---------------------------------------+ - --- SST index search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -DROP TABLE vectors_high_dim; - -Affected Rows: 0 - --- ============================================ --- Part 8: Different k values (LIMIT) --- ============================================ -CREATE TABLE vectors_k ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_k VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.8, 0.2, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.7, 0.3, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.6, 0.4, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.5, 0.5, 0.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.4, 0.6, 0.0, 0.0]'), - (8, '2024-01-01 00:00:07', '[0.3, 0.7, 0.0, 0.0]'), - (9, '2024-01-01 00:00:08', '[0.2, 0.8, 0.0, 0.0]'), - (10, '2024-01-01 00:00:09', '[0.1, 0.9, 0.0, 0.0]'); - -Affected Rows: 10 - -ADMIN FLUSH_TABLE('vectors_k'); - -+--------------------------------+ -| ADMIN FLUSH_TABLE('vectors_k') | -+--------------------------------+ -| 0 | -+--------------------------------+ - --- k=1 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 1; - -+--------+ -| vec_id | -+--------+ -| 1 | -+--------+ - --- k=5 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 5; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 2 | -| 3 | -| 4 | -| 5 | -+--------+ - --- k=10 (all vectors) -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 10; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 2 | -| 3 | -| 4 | -| 5 | -| 6 | -| 7 | -| 8 | -| 9 | -| 10 | -+--------+ - -DROP TABLE vectors_k; - -Affected Rows: 0 - --- ============================================ --- Part 9: Engine parameter tests --- ============================================ --- Create table with explicit engine parameter (usearch is default) -CREATE TABLE vectors_engine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (engine = 'usearch', metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert test vectors -INSERT INTO vectors_engine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0]'); - -Affected Rows: 3 - --- Memtable search -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -ADMIN FLUSH_TABLE('vectors_engine'); - -+-------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_engine') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SST index search with usearch engine -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -DROP TABLE vectors_engine; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/vector/vector_index.sql b/tests/cases/standalone/common/function/vector/vector_index.sql deleted file mode 100644 index bbcf02d7f30..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index.sql +++ /dev/null @@ -1,341 +0,0 @@ --- Test vector index creation and KNN search - --- ============================================ --- Part 1: Basic L2 squared distance tests --- ============================================ - --- Create a table with vector column and L2sq vector index -CREATE TABLE vectors_l2sq ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert test vectors -INSERT INTO vectors_l2sq VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.0, 0.0, 1.0]'), - (5, '2024-01-01 00:00:04', '[1.0, 1.0, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.0, 1.0, 1.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.0, 0.0, 1.0, 1.0]'), - (8, '2024-01-01 00:00:07', '[1.0, 0.0, 0.0, 1.0]'); - --- Query BEFORE flush (memtable search) --- Expected: vec_id=1 (distance=0), vec_id=5 (distance=1), vec_id=8 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - --- Flush to create SST files with vector index -ADMIN FLUSH_TABLE('vectors_l2sq'); - --- Query AFTER flush (SST index search) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - --- Query with different target vector --- Expected: vec_id=6 (distance=0), vec_id=2 (distance=1), vec_id=3 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_l2sq; - --- ============================================ --- Part 2: Cosine distance tests --- ============================================ - -CREATE TABLE vectors_cosine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - --- Insert vectors with different magnitudes but same/different directions -INSERT INTO vectors_cosine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - --- Memtable search with cosine distance --- vec_id=1 and vec_id=2 should have same cosine distance (0) since they point same direction -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -ADMIN FLUSH_TABLE('vectors_cosine'); - --- SST index search with cosine distance -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_cosine; - --- ============================================ --- Part 3: Dot product (inner product) tests --- ============================================ - -CREATE TABLE vectors_dot ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'dot'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_dot VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - --- Memtable search with dot product --- Larger dot product means more similar, so we use negative for ordering --- vec_id=2 should be best (dot=2), vec_id=1 and vec_id=4 have dot=1 -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -ADMIN FLUSH_TABLE('vectors_dot'); - --- SST index search with dot product -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -DROP TABLE vectors_dot; - --- ============================================ --- Part 4: NULL vector handling tests --- ============================================ - -CREATE TABLE vectors_null ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert vectors with some NULLs -INSERT INTO vectors_null VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', NULL), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', NULL), - (5, '2024-01-01 00:00:04', '[0.0, 0.0, 1.0, 0.0]'); - --- Memtable search should skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -ADMIN FLUSH_TABLE('vectors_null'); - --- SST index search should also skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_null; - --- ============================================ --- Part 5: Mixed memtable + SST search tests --- ============================================ - -CREATE TABLE vectors_mixed ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert first batch and flush to SST -INSERT INTO vectors_mixed VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'); - -ADMIN FLUSH_TABLE('vectors_mixed'); - --- Insert second batch (stays in memtable) -INSERT INTO vectors_mixed VALUES - (4, '2024-01-01 00:00:03', '[0.5, 0.5, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.9, 0.1, 0.0, 0.0]'); - --- Query should search both SST (vec_id 1,2,3) and memtable (vec_id 4,5) --- Target: [1.0, 0.0, 0.0, 0.0] --- Expected: vec_id=1 (dist=0), vec_id=5 (dist=0.02), vec_id=4 (dist=0.5) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_mixed -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_mixed; - --- ============================================ --- Part 6: KNN with WHERE clause tests --- ============================================ - -CREATE TABLE vectors_filter ( - vec_id INT, - "category" STRING, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_filter VALUES - (1, 'A', '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, 'B', '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, 'A', '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, 'B', '2024-01-01 00:00:03', '[0.1, 0.9, 0.0, 0.0]'), - (5, 'A', '2024-01-01 00:00:04', '[0.5, 0.5, 0.0, 0.0]'); - --- Memtable search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -ADMIN FLUSH_TABLE('vectors_filter'); - --- SST index search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - --- Filter with time range -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE ts >= '2024-01-01 00:00:02' AND ts <= '2024-01-01 00:00:04' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_filter; - --- ============================================ --- Part 7: Higher dimension vectors --- ============================================ - -CREATE TABLE vectors_high_dim ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(128) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert high-dimensional vectors (simplified: first few elements differ) -INSERT INTO vectors_high_dim VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'); - --- Memtable search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -ADMIN FLUSH_TABLE('vectors_high_dim'); - --- SST index search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_high_dim; - --- ============================================ --- Part 8: Different k values (LIMIT) --- ============================================ - -CREATE TABLE vectors_k ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_k VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.8, 0.2, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.7, 0.3, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.6, 0.4, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.5, 0.5, 0.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.4, 0.6, 0.0, 0.0]'), - (8, '2024-01-01 00:00:07', '[0.3, 0.7, 0.0, 0.0]'), - (9, '2024-01-01 00:00:08', '[0.2, 0.8, 0.0, 0.0]'), - (10, '2024-01-01 00:00:09', '[0.1, 0.9, 0.0, 0.0]'); - -ADMIN FLUSH_TABLE('vectors_k'); - --- k=1 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 1; - --- k=5 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 5; - --- k=10 (all vectors) -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 10; - -DROP TABLE vectors_k; - --- ============================================ --- Part 9: Engine parameter tests --- ============================================ - --- Create table with explicit engine parameter (usearch is default) -CREATE TABLE vectors_engine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (engine = 'usearch', metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert test vectors -INSERT INTO vectors_engine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0]'); - --- Memtable search -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -ADMIN FLUSH_TABLE('vectors_engine'); - --- SST index search with usearch engine -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_engine; diff --git a/tests/cases/standalone/common/function/vector/vector_index_explain.result b/tests/cases/standalone/common/function/vector/vector_index_explain.result deleted file mode 100644 index 60f01ce157e..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_explain.result +++ /dev/null @@ -1,272 +0,0 @@ --- Vector index explain analyze coverage --- ============================================ --- Part 1: Single table KNN explain --- ============================================ -CREATE TABLE vectors_explain ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_explain VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -ADMIN FLUSH_TABLE('vectors_explain'); - -+--------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_explain') | -+--------------------------------------+ -| 0 | -+--------------------------------------+ - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 ,\s"dyn_filters":\s\[.* REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_CooperativeExec metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], fetch=2 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=2), expr=[vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["vec_id", "embedding"]REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - -DROP TABLE vectors_explain; - -Affected Rows: 0 - --- ============================================ --- Part 2: Join with vector order/limit --- ============================================ -CREATE TABLE vectors_explain_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -CREATE TABLE vectors_explain_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_explain_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -INSERT INTO vectors_explain_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -Affected Rows: 2 - -ADMIN FLUSH_TABLE('vectors_explain_left'); - -+-------------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_explain_left') | -+-------------------------------------------+ -| 0 | -+-------------------------------------------+ - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT l.vec_id -FROM vectors_explain_left l -JOIN vectors_explain_right r ON l.vec_id = r.vec_id -ORDER BY vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), l.vec_id -LIMIT 1; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], fetch=1 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=1), expr=[vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(vec_id@0, vec_id@0)], projection=[vec_id@0, embedding@1] metrics=REDACTED_| -|_|_|_RepartitionExec: partitioning=Hash([vec_id@0],REDACTED -|_|_|_ProjectionExec: expr=[vec_id@0 as vec_id, embedding@2 as embedding] metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_RepartitionExec: partitioning=Hash([vec_id@0],REDACTED -|_|_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_FilterExec: vec_id@0 IS NOT NULL metrics=REDACTED_| -|_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -| 1_| 0_|_FilterExec: vec_id@0 IS NOT NULL metrics=REDACTED_| -|_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - -DROP TABLE vectors_explain_left; - -Affected Rows: 0 - -DROP TABLE vectors_explain_right; - -Affected Rows: 0 - --- ============================================ --- Part 3: Cosine and dot explain coverage --- ============================================ -CREATE TABLE vectors_explain_metric ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_explain_metric VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0]'), - (3, '2024-01-01 00:00:02', '[-1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, -1.0]'); - -Affected Rows: 4 - -ADMIN FLUSH_TABLE('vectors_explain_metric'); - -+---------------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_explain_metric') | -+---------------------------------------------+ -| 0 | -+---------------------------------------------+ - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 SeqScan:.* SeqScan: REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_CooperativeExec metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_cos_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], fetch=2 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=2), expr=[vec_cos_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_dot_product(embedding, '[1.0, 0.0]') DESC, vec_id -LIMIT 2; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_CooperativeExec metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_dot_product(embedding@1, [1.0, 0.0]) DESC, vec_id@0 ASC NULLS LAST], fetch=2 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=2), expr=[vec_dot_product(embedding@1, [1.0, 0.0]) DESC, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - -DROP TABLE vectors_explain_metric; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/vector/vector_index_explain.sql b/tests/cases/standalone/common/function/vector/vector_index_explain.sql deleted file mode 100644 index 7a0330b0ea2..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_explain.sql +++ /dev/null @@ -1,163 +0,0 @@ --- Vector index explain analyze coverage - - --- ============================================ --- Part 1: Single table KNN explain --- ============================================ -CREATE TABLE vectors_explain ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_explain VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -ADMIN FLUSH_TABLE('vectors_explain'); - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 ,\s"dyn_filters":\s\[.* REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_explain; - --- ============================================ --- Part 2: Join with vector order/limit --- ============================================ -CREATE TABLE vectors_explain_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -CREATE TABLE vectors_explain_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_explain_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -INSERT INTO vectors_explain_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -ADMIN FLUSH_TABLE('vectors_explain_left'); - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT l.vec_id -FROM vectors_explain_left l -JOIN vectors_explain_right r ON l.vec_id = r.vec_id -ORDER BY vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), l.vec_id -LIMIT 1; - -DROP TABLE vectors_explain_left; -DROP TABLE vectors_explain_right; - --- ============================================ --- Part 3: Cosine and dot explain coverage --- ============================================ -CREATE TABLE vectors_explain_metric ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_explain_metric VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0]'), - (3, '2024-01-01 00:00:02', '[-1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, -1.0]'); - -ADMIN FLUSH_TABLE('vectors_explain_metric'); - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 SeqScan:.* SeqScan: REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - --- 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\(\[vec_id@0\],.* Hash([vec_id@0],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 SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_dot_product(embedding, '[1.0, 0.0]') DESC, vec_id -LIMIT 2; - -DROP TABLE vectors_explain_metric; diff --git a/tests/cases/standalone/common/function/vector/vector_index_join.result b/tests/cases/standalone/common/function/vector/vector_index_join.result deleted file mode 100644 index 51fb3934058..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_join.result +++ /dev/null @@ -1,240 +0,0 @@ --- Vector index join/subquery coverage --- ============================================ --- Part 1: Join should not pre-limit left table --- ============================================ -CREATE TABLE vectors_join_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -CREATE TABLE vectors_join_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_join_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -INSERT INTO vectors_join_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -Affected Rows: 2 - -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_join_left l -JOIN vectors_join_right r ON l.vec_id = r.vec_id -ORDER BY dist, l.vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_join_left; - -Affected Rows: 0 - -DROP TABLE vectors_join_right; - -Affected Rows: 0 - --- ============================================ --- Part 2: Subquery should be a barrier --- ============================================ -CREATE TABLE vectors_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_subquery WHERE vec_id >= 3 -) s -ORDER BY dist, s.vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_subquery; - -Affected Rows: 0 - --- ============================================ --- Part 3: LEFT JOIN should not pre-limit --- ============================================ -CREATE TABLE vectors_left_join ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -CREATE TABLE vectors_left_join_filter ( - vec_id INT, - keep BOOLEAN, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_left_join VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - --- Only vec_id 3,4 have matching rows in filter table -INSERT INTO vectors_left_join_filter VALUES - (3, true, '2024-01-01 00:00:02'), - (4, true, '2024-01-01 00:00:03'); - -Affected Rows: 2 - --- LEFT JOIN then filter by IS NOT NULL --- Should return vec_id=4 (dist=1.81), not vec_id=1 or 2 -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_left_join l -LEFT JOIN vectors_left_join_filter r ON l.vec_id = r.vec_id -WHERE r.vec_id IS NOT NULL -ORDER BY dist, l.vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_left_join; - -Affected Rows: 0 - -DROP TABLE vectors_left_join_filter; - -Affected Rows: 0 - --- ============================================ --- Part 4: Inlineable subquery should allow hint --- ============================================ -CREATE TABLE vectors_inline_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_inline_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -ADMIN FLUSH_TABLE('vectors_inline_subquery'); - -+----------------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_inline_subquery') | -+----------------------------------------------+ -| 0 | -+----------------------------------------------+ - --- Subquery without LIMIT/DISTINCT/aggregation can be inlined --- Vector hint should be able to push down -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_inline_subquery WHERE vec_id >= 1 -) s -ORDER BY dist, s.vec_id -LIMIT 2; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 1 | 0.0 | -| 2 | 0.01 | -+--------+------+ - -DROP TABLE vectors_inline_subquery; - -Affected Rows: 0 - --- ============================================ --- Part 5: CTE should be a barrier --- ============================================ -CREATE TABLE vectors_cte ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_cte VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - --- CTE acts as optimization barrier --- Filter in CTE limits to vec_id >= 3, so result should be vec_id=4 (dist=1.81) -WITH filtered AS ( - SELECT * FROM vectors_cte WHERE vec_id >= 3 -) -SELECT vec_id, round(vec_l2sq_distance(embedding, '[1.0, 0.0]'), 2) AS dist -FROM filtered -ORDER BY dist, vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_cte; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/vector/vector_index_join.sql b/tests/cases/standalone/common/function/vector/vector_index_join.sql deleted file mode 100644 index 699942b3001..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_join.sql +++ /dev/null @@ -1,159 +0,0 @@ --- Vector index join/subquery coverage - --- ============================================ --- Part 1: Join should not pre-limit left table --- ============================================ -CREATE TABLE vectors_join_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -CREATE TABLE vectors_join_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_join_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -INSERT INTO vectors_join_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_join_left l -JOIN vectors_join_right r ON l.vec_id = r.vec_id -ORDER BY dist, l.vec_id -LIMIT 1; - -DROP TABLE vectors_join_left; -DROP TABLE vectors_join_right; - --- ============================================ --- Part 2: Subquery should be a barrier --- ============================================ -CREATE TABLE vectors_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_subquery WHERE vec_id >= 3 -) s -ORDER BY dist, s.vec_id -LIMIT 1; - -DROP TABLE vectors_subquery; - --- ============================================ --- Part 3: LEFT JOIN should not pre-limit --- ============================================ -CREATE TABLE vectors_left_join ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -CREATE TABLE vectors_left_join_filter ( - vec_id INT, - keep BOOLEAN, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_left_join VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - --- Only vec_id 3,4 have matching rows in filter table -INSERT INTO vectors_left_join_filter VALUES - (3, true, '2024-01-01 00:00:02'), - (4, true, '2024-01-01 00:00:03'); - --- LEFT JOIN then filter by IS NOT NULL --- Should return vec_id=4 (dist=1.81), not vec_id=1 or 2 -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_left_join l -LEFT JOIN vectors_left_join_filter r ON l.vec_id = r.vec_id -WHERE r.vec_id IS NOT NULL -ORDER BY dist, l.vec_id -LIMIT 1; - -DROP TABLE vectors_left_join; -DROP TABLE vectors_left_join_filter; - --- ============================================ --- Part 4: Inlineable subquery should allow hint --- ============================================ -CREATE TABLE vectors_inline_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_inline_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -ADMIN FLUSH_TABLE('vectors_inline_subquery'); - --- Subquery without LIMIT/DISTINCT/aggregation can be inlined --- Vector hint should be able to push down -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_inline_subquery WHERE vec_id >= 1 -) s -ORDER BY dist, s.vec_id -LIMIT 2; - -DROP TABLE vectors_inline_subquery; - --- ============================================ --- Part 5: CTE should be a barrier --- ============================================ -CREATE TABLE vectors_cte ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_cte VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - --- CTE acts as optimization barrier --- Filter in CTE limits to vec_id >= 3, so result should be vec_id=4 (dist=1.81) -WITH filtered AS ( - SELECT * FROM vectors_cte WHERE vec_id >= 3 -) -SELECT vec_id, round(vec_l2sq_distance(embedding, '[1.0, 0.0]'), 2) AS dist -FROM filtered -ORDER BY dist, vec_id -LIMIT 1; - -DROP TABLE vectors_cte; diff --git a/tests/cases/standalone/common/information_schema/region_statistics.result b/tests/cases/standalone/common/information_schema/region_statistics.result index e6f9c9e47fc..5589b5d00e2 100644 --- a/tests/cases/standalone/common/information_schema/region_statistics.result +++ b/tests/cases/standalone/common/information_schema/region_statistics.result @@ -21,6 +21,34 @@ INSERT INTO test VALUES Affected Rows: 3 +-- The projected columns are millisecond-typed. A nanosecond time index proves +-- min_timestamp floors and max_timestamp ceils, so the reported window stays a +-- superset of the data and cannot hide a region from a time-bounded lookup. +CREATE TABLE precise ( + a int primary key, + ts timestamp(9) time index, +); + +Affected Rows: 0 + +-- 1500000ns = 1.5ms floors to 1ms; 3999999ns = 3.999999ms ceils to 4ms. +INSERT INTO precise VALUES + (1, 1500000), + (2, 2500000), + (3, 3999999); + +Affected Rows: 3 + +-- A region with no rows has nothing to bound, so both ends stay NULL rather +-- than collapsing onto the epoch. +CREATE TABLE empty_region ( + a int primary key, + ts timestamp time index, +); + +Affected Rows: 0 + +-- Wait once for the region statistics of all three tables to be reported. -- SQLNESS SLEEP 3s -- SQLNESS REPLACE (\s+\d+\s+) -- For regions using different WAL implementations, the manifest size may vary. @@ -49,25 +77,6 @@ DROP TABLE test; Affected Rows: 0 --- The projected columns are millisecond-typed. A nanosecond time index proves --- min_timestamp floors and max_timestamp ceils, so the reported window stays a --- superset of the data and cannot hide a region from a time-bounded lookup. -CREATE TABLE precise ( - a int primary key, - ts timestamp(9) time index, -); - -Affected Rows: 0 - --- 1500000ns = 1.5ms floors to 1ms; 3999999ns = 3.999999ms ceils to 4ms. -INSERT INTO precise VALUES - (1, 1500000), - (2, 2500000), - (3, 3999999); - -Affected Rows: 3 - --- SQLNESS SLEEP 3s SELECT min_timestamp, max_timestamp FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); @@ -82,16 +91,6 @@ DROP TABLE precise; Affected Rows: 0 --- A region with no rows has nothing to bound, so both ends stay NULL rather --- than collapsing onto the epoch. -CREATE TABLE empty_region ( - a int primary key, - ts timestamp time index, -); - -Affected Rows: 0 - --- SQLNESS SLEEP 3s SELECT min_timestamp, max_timestamp FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'empty_region' and table_schema = 'public'); diff --git a/tests/cases/standalone/common/information_schema/region_statistics.sql b/tests/cases/standalone/common/information_schema/region_statistics.sql index 60264fbfb57..6ce12b96417 100644 --- a/tests/cases/standalone/common/information_schema/region_statistics.sql +++ b/tests/cases/standalone/common/information_schema/region_statistics.sql @@ -16,20 +16,6 @@ INSERT INTO test VALUES (11, 'b', 11), (21, 'c', 21); --- SQLNESS SLEEP 3s --- SQLNESS REPLACE (\s+\d+\s+) --- For regions using different WAL implementations, the manifest size may vary. --- The remote WAL implementation additionally stores a flushed entry ID when creating the manifest. -SELECT SUM(region_rows), SUM(written_bytes_since_open), SUM(query_cpu_time_millis), - SUM(query_scanned_bytes), SUM(memtable_size), SUM(sst_size), SUM(index_size) - FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id - IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test' and table_schema = 'public'); - --- SQLNESS REPLACE (\s+\d+\s+) -SELECT data_length, index_length, avg_row_length, table_rows FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test'; - -DROP TABLE test; - -- The projected columns are millisecond-typed. A nanosecond time index proves -- min_timestamp floors and max_timestamp ceils, so the reported window stays a -- superset of the data and cannot hide a region from a time-bounded lookup. @@ -44,13 +30,6 @@ INSERT INTO precise VALUES (2, 2500000), (3, 3999999); --- SQLNESS SLEEP 3s -SELECT min_timestamp, max_timestamp - FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id - IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); - -DROP TABLE precise; - -- A region with no rows has nothing to bound, so both ends stay NULL rather -- than collapsing onto the epoch. CREATE TABLE empty_region ( @@ -58,7 +37,27 @@ CREATE TABLE empty_region ( ts timestamp time index, ); +-- Wait once for the region statistics of all three tables to be reported. -- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (\s+\d+\s+) +-- For regions using different WAL implementations, the manifest size may vary. +-- The remote WAL implementation additionally stores a flushed entry ID when creating the manifest. +SELECT SUM(region_rows), SUM(written_bytes_since_open), SUM(query_cpu_time_millis), + SUM(query_scanned_bytes), SUM(memtable_size), SUM(sst_size), SUM(index_size) + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test' and table_schema = 'public'); + +-- SQLNESS REPLACE (\s+\d+\s+) +SELECT data_length, index_length, avg_row_length, table_rows FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test'; + +DROP TABLE test; + +SELECT min_timestamp, max_timestamp + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); + +DROP TABLE precise; + SELECT min_timestamp, max_timestamp FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'empty_region' and table_schema = 'public'); diff --git a/tests/cases/standalone/common/information_schema/ssts.result b/tests/cases/standalone/common/information_schema/ssts.result index 9777e3b12a9..289b4599387 100644 --- a/tests/cases/standalone/common/information_schema/ssts.result +++ b/tests/cases/standalone/common/information_schema/ssts.result @@ -229,6 +229,20 @@ SELECT * FROM information_schema.ssts_storage order by file_path; | data/greptime/public//_/index/.puffin || || +---------------------------------------------------------------------------------------------+-----------+-------------------------+---------+ +-- Regression: repeated physical optimization must preserve the global DISTINCT limit. +SELECT COUNT(*) AS distinct_limited_rows +FROM ( + SELECT DISTINCT region_id + FROM information_schema.ssts_manifest + LIMIT 1 +); + ++-----------------------+ +| distinct_limited_rows | ++-----------------------+ +| 1 | ++-----------------------+ + DROP TABLE sst_case; Affected Rows: 0 diff --git a/tests/cases/standalone/common/information_schema/ssts.sql b/tests/cases/standalone/common/information_schema/ssts.sql index 43512de6808..7b40b0056bb 100644 --- a/tests/cases/standalone/common/information_schema/ssts.sql +++ b/tests/cases/standalone/common/information_schema/ssts.sql @@ -72,4 +72,12 @@ SELECT * FROM information_schema.ssts_index_meta ORDER BY meta_json; -- SQLNESS REPLACE (/public/\d+/\d+_\d+) /public//_ SELECT * FROM information_schema.ssts_storage order by file_path; +-- Regression: repeated physical optimization must preserve the global DISTINCT limit. +SELECT COUNT(*) AS distinct_limited_rows +FROM ( + SELECT DISTINCT region_id + FROM information_schema.ssts_manifest + LIMIT 1 +); + DROP TABLE sst_case; diff --git a/tests/cases/standalone/common/promql/at_modifier.result b/tests/cases/standalone/common/promql/at_modifier.result new file mode 100644 index 00000000000..27570231833 --- /dev/null +++ b/tests/cases/standalone/common/promql/at_modifier.result @@ -0,0 +1,750 @@ +-- Tests for the PromQL `@` modifier on vector and matrix selectors. +-- +-- `@` anchors the sample selection window at a fixed timestamp instead of the timestamp of each +-- evaluation step: `@ ` uses the given timestamp, `@ start()` / `@ end()` use the +-- start/end of the statement's evaluation range, and `offset` shifts the anchor backwards before +-- the window is selected. The output timestamps still follow the evaluation grid. +-- +-- Every metric below carries two series, host 'a' and host 'b', with different values: the +-- anchored results must therefore report both of them, per step. A series silently dropped, or a +-- batch holding several series handled as one timeline, does not look like a correct single-series +-- answer here but shows up as a missing or shifted row. +-- +-- Sample timestamps below are in milliseconds, while `@` and `TQL EVAL` timestamps are in seconds, +-- as in Prometheus. +CREATE TABLE at_modifier_gauge ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +CREATE TABLE at_modifier_counter_total ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +-- One sample per minute, so the anchored sample is easy to identify. The default lookback is 5m. +-- Host 'a' reports 0.0 to 6.0, host 'b' reports 7.0 to 13.0: the two series are told apart by their +-- values alone, which stays true after any anchoring. +INSERT INTO at_modifier_gauge VALUES + (0, 0.0, 'a'), + (60000, 1.0, 'a'), + (120000, 2.0, 'a'), + (180000, 3.0, 'a'), + (240000, 4.0, 'a'), + (300000, 5.0, 'a'), + (360000, 6.0, 'a'), + (0, 7.0, 'b'), + (60000, 8.0, 'b'), + (120000, 9.0, 'b'), + (180000, 10.0, 'b'), + (240000, 11.0, 'b'), + (300000, 12.0, 'b'), + (360000, 13.0, 'b'); + +Affected Rows: 14 + +-- Host 'a': a counter increasing by 60 per minute, i.e. by 1 per second. +-- Host 'b': a counter starting at 100 and increasing by 30 per minute, i.e. by 0.5 per second. Its +-- rate differs from host 'a', so a window folded for the wrong series cannot pass as the right one. +INSERT INTO at_modifier_counter_total VALUES + (0, 0.0, 'a'), + (60000, 60.0, 'a'), + (120000, 120.0, 'a'), + (180000, 180.0, 'a'), + (240000, 240.0, 'a'), + (300000, 300.0, 'a'), + (360000, 360.0, 'a'), + (0, 100.0, 'b'), + (60000, 130.0, 'b'), + (120000, 160.0, 'b'), + (180000, 190.0, 'b'), + (240000, 220.0, 'b'), + (300000, 250.0, 'b'), + (360000, 280.0, 'b'); + +Affected Rows: 14 + +-- 1. Instant query anchored at an absolute timestamp: the sample window is `(anchor - lookback, +-- anchor]` = `(0s, 300s]`, so its newest samples are 5.0 for host 'a' and 12.0 for host 'b', both +-- at 300s. The output timestamps stay the evaluation timestamps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 12.0 | b | +| 1970-01-01T00:01:40 | 5.0 | a | ++---------------------+------+------+ + +-- Control: the same query without `@` looks back from the evaluation timestamp instead: +-- `(-200s, 100s]` selects the samples 1.0 (a) and 8.0 (b) at 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge; + ++---------------------+-----+------+ +| ts | val | host | ++---------------------+-----+------+ +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 8.0 | b | ++---------------------+-----+------+ + +-- The anchor, not the evaluation timestamp, decides which sample is reported. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 360; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 13.0 | b | +| 1970-01-01T00:01:40 | 6.0 | a | ++---------------------+------+------+ + +-- Sub-second precision is preserved in the anchor: `240.999` selects the samples at 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 240.999; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 11.0 | b | +| 1970-01-01T00:01:40 | 4.0 | a | ++---------------------+------+------+ + +-- The window `(-300s, 0s]` includes the samples exactly at the anchor. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 0; + ++---------------------+-----+------+ +| ts | val | host | ++---------------------+-----+------+ +| 1970-01-01T00:01:40 | 0.0 | a | +| 1970-01-01T00:01:40 | 7.0 | b | ++---------------------+-----+------+ + +-- 2. `@ start()`: every step selects its samples around the start of the evaluation range +-- (`(-200s, 100s]`), so every step reports the same value per series (1.0 for 'a', 8.0 for 'b'). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ start(); + ++---------------------+-----+------+ +| ts | val | host | ++---------------------+-----+------+ +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 8.0 | b | +| 1970-01-01T00:03:20 | 1.0 | a | +| 1970-01-01T00:03:20 | 8.0 | b | +| 1970-01-01T00:05:00 | 1.0 | a | +| 1970-01-01T00:05:00 | 8.0 | b | +| 1970-01-01T00:06:40 | 1.0 | a | +| 1970-01-01T00:06:40 | 8.0 | b | ++---------------------+-----+------+ + +-- Control: without `@` every step looks back from its own timestamp, so the steps differ. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 8.0 | b | +| 1970-01-01T00:03:20 | 10.0 | b | +| 1970-01-01T00:03:20 | 3.0 | a | +| 1970-01-01T00:05:00 | 12.0 | b | +| 1970-01-01T00:05:00 | 5.0 | a | +| 1970-01-01T00:06:40 | 13.0 | b | +| 1970-01-01T00:06:40 | 6.0 | a | ++---------------------+------+------+ + +-- 3. `@ end()`: every step selects its samples around the end of the evaluation range +-- (`(100s, 400s]`), whose newest samples are 6.0 (a) and 13.0 (b) at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ end(); + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 13.0 | b | +| 1970-01-01T00:01:40 | 6.0 | a | +| 1970-01-01T00:03:20 | 13.0 | b | +| 1970-01-01T00:03:20 | 6.0 | a | +| 1970-01-01T00:05:00 | 13.0 | b | +| 1970-01-01T00:05:00 | 6.0 | a | +| 1970-01-01T00:06:40 | 13.0 | b | +| 1970-01-01T00:06:40 | 6.0 | a | ++---------------------+------+------+ + +-- 4. Range selector anchored by `@`: the call directly above the anchored range selector is +-- evaluated once, at the start of the evaluation, and its result is reported at every step — the +-- planner's counterpart of Prometheus' `StepInvariantExpr`. The window `(0s, 300s]` is left-open, +-- so it holds 60 -> 300 for host 'a' and 130 -> 250 for host 'b', i.e. 240 and 120 units. The 60s +-- gap to the window start is below the extrapolation threshold (1.1 x the 60s sampling interval), +-- so `rate` adds it and divides by the 300s window: 240 * 1.25 / 300 = 1.0 and +-- 120 * 1.25 / 300 = 0.5 per second. Every step must report those same per-series values, and both +-- series must be there at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 1.0 | a | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 1.0 | a | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 1.0 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 1.0 | a | ++---------------------+------------------------------------------+------+ + +-- Consistency: the instant query below shows the anchored values over the same left-open +-- `(0s, 300s]` window (1.0 per second for 'a', 0.5 for 'b'). The `!= 1.0` comparison filters +-- host 'a', but host 'b' remains at every step with its rate of 0.5 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 300, '1s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:05:00 | 0.5 | b | +| 1970-01-01T00:05:00 | 1.0 | a | ++---------------------+------------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300) != 1.0; + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.5 | b | ++---------------------+------------------------------------------+------+ + +-- Parentheses around the range argument are transparent: the call must be promoted exactly like +-- the unparenthesized one above, so every step reports the same anchored values for both series +-- (1.0 for host 'a', 0.5 for host 'b') instead of folding a window per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate((at_modifier_counter_total[5m] @ 300)); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 1.0 | a | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 1.0 | a | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 1.0 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 1.0 | a | ++---------------------+------------------------------------------+------+ + +-- The same holds for `@ start()` and `@ end()`, which are also fixed anchors: they resolve to the +-- statement's evaluation range, so a range call using them is evaluated once as well. +-- The statement starts at 60s: the window `(-240s, 60s]` holds only the samples at 0s and 60s, so +-- host 'a' goes 0 -> 60 and host 'b' 100 -> 130. The far 240s gap to the window start is beyond the +-- extrapolation threshold (1.1 x the 60s interval), so only half an interval (30s) is added there; +-- host 'a' starts its counter at 0, which clamps that back: 60 / 300 = 0.2 and +-- 30 * 1.5 / 300 = 0.15 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 300, '60s') rate(at_modifier_counter_total[5m] @ start()); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:01:00 | 0.15 | b | +| 1970-01-01T00:01:00 | 0.2 | a | +| 1970-01-01T00:02:00 | 0.15 | b | +| 1970-01-01T00:02:00 | 0.2 | a | +| 1970-01-01T00:03:00 | 0.15 | b | +| 1970-01-01T00:03:00 | 0.2 | a | +| 1970-01-01T00:04:00 | 0.15 | b | +| 1970-01-01T00:04:00 | 0.2 | a | +| 1970-01-01T00:05:00 | 0.15 | b | +| 1970-01-01T00:05:00 | 0.2 | a | ++---------------------+------------------------------------------+------+ + +-- The instant query at 60s folds the same window as every step of the query above. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 60, '1s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:01:00 | 0.15 | b | +| 1970-01-01T00:01:00 | 0.2 | a | ++---------------------+------------------------------------------+------+ + +-- The statement ends at 240s: the window `(-60s, 240s]` holds the samples 0s to 240s, so host 'a' +-- goes 0 -> 240 and host 'b' 100 -> 220. The 60s gap to the window start is added for both, but +-- host 'a' starts at zero, which clamps it back (240 / 300 = 0.8 per second) while host 'b' keeps +-- it (120 * 1.25 / 300 = 0.5 per second). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ end()); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:00:00 | 0.8 | a | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 0.8 | a | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 0.8 | a | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 0.8 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.8 | a | ++---------------------+------------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (240, 240, '1s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.8 | a | ++---------------------+------------------------------------------+------+ + +-- Control: without `@` each step folds its own window, so the values keep changing. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:01:00 | 0.15 | b | +| 1970-01-01T00:01:00 | 0.2 | a | +| 1970-01-01T00:02:00 | 0.25 | b | +| 1970-01-01T00:02:00 | 0.4 | a | +| 1970-01-01T00:03:00 | 0.35000000000000003 | b | +| 1970-01-01T00:03:00 | 0.6000000000000001 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.8 | a | ++---------------------+------------------------------------------+------+ + +-- `@ start()` on a range selector: the window `(-300s, 0s]` is left-open, so it holds a single +-- sample per series — the one exactly at the anchor — and `count_over_time` returns 1 at every step +-- for both of them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m] @ start()); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 1.0 | b | +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 1.0 | b | +| 1970-01-01T00:03:20 | 1.0 | a | +| 1970-01-01T00:03:20 | 1.0 | b | +| 1970-01-01T00:05:00 | 1.0 | a | +| 1970-01-01T00:05:00 | 1.0 | b | ++---------------------+------------------------------------+------+ + +-- Control: without `@` the window keeps moving, so the count grows with the step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 1.0 | b | +| 1970-01-01T00:01:40 | 2.0 | a | +| 1970-01-01T00:01:40 | 2.0 | b | +| 1970-01-01T00:03:20 | 4.0 | a | +| 1970-01-01T00:03:20 | 4.0 | b | +| 1970-01-01T00:05:00 | 5.0 | a | +| 1970-01-01T00:05:00 | 5.0 | b | ++---------------------+------------------------------------+------+ + +-- 5. `@` combined with `offset`: the offset moves the anchor backwards before the window is +-- selected, so `@ 300 offset 2m` reports the samples at 180s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 2m; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 10.0 | b | +| 1970-01-01T00:01:40 | 3.0 | a | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 1m; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 11.0 | b | +| 1970-01-01T00:01:40 | 4.0 | a | ++---------------------+------+------+ + +-- Control: the offset applies to the evaluation timestamp when there is no `@`. +TQL EVAL (100, 100, '1s') at_modifier_gauge offset 2m; + ++----+-----+------+ +| ts | val | host | ++----+-----+------+ ++----+-----+------+ + +-- 6. `@` inside a larger expression. The anchored operand contributes the same value at every +-- step, while the plain operand still follows the evaluation step. A binary expression is never +-- promoted: when one side is anchored and the other is not, the anchored side replays on its own, +-- and when both sides are anchored, each side reports its own anchored samples. The join then runs +-- at every step over the replayed per-series inputs. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge @ 0; + ++------+---------------------+-------------------+ +| host | ts | lhs.val + rhs.val | ++------+---------------------+-------------------+ +| a | 1970-01-01T00:01:40 | 5.0 | +| a | 1970-01-01T00:03:20 | 5.0 | +| a | 1970-01-01T00:05:00 | 5.0 | +| a | 1970-01-01T00:06:40 | 5.0 | +| b | 1970-01-01T00:01:40 | 19.0 | +| b | 1970-01-01T00:03:20 | 19.0 | +| b | 1970-01-01T00:05:00 | 19.0 | +| b | 1970-01-01T00:06:40 | 19.0 | ++------+---------------------+-------------------+ + +-- The aggregation above the anchored selector runs at every step over the replayed (still +-- per-series) samples, so it sums both series at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') sum(at_modifier_gauge @ 0); + ++---------------------+----------------------------+ +| ts | sum(at_modifier_gauge.val) | ++---------------------+----------------------------+ +| 1970-01-01T00:01:40 | 7.0 | +| 1970-01-01T00:03:20 | 7.0 | +| 1970-01-01T00:05:00 | 7.0 | +| 1970-01-01T00:06:40 | 7.0 | ++---------------------+----------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 0); + ++---------------------+----------+------+ +| ts | abs(val) | host | ++---------------------+----------+------+ +| 1970-01-01T00:01:40 | 0.0 | a | +| 1970-01-01T00:01:40 | 7.0 | b | +| 1970-01-01T00:03:20 | 0.0 | a | +| 1970-01-01T00:03:20 | 7.0 | b | +| 1970-01-01T00:05:00 | 0.0 | a | +| 1970-01-01T00:05:00 | 7.0 | b | +| 1970-01-01T00:06:40 | 0.0 | a | +| 1970-01-01T00:06:40 | 7.0 | b | ++---------------------+----------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge; + ++------+---------------------+-------------------+ +| host | ts | lhs.val + rhs.val | ++------+---------------------+-------------------+ +| a | 1970-01-01T00:01:40 | 6.0 | +| a | 1970-01-01T00:03:20 | 8.0 | +| a | 1970-01-01T00:05:00 | 10.0 | +| a | 1970-01-01T00:06:40 | 11.0 | +| b | 1970-01-01T00:01:40 | 20.0 | +| b | 1970-01-01T00:03:20 | 22.0 | +| b | 1970-01-01T00:05:00 | 24.0 | +| b | 1970-01-01T00:06:40 | 25.0 | ++------+---------------------+-------------------+ + +-- 7. `@ -1`, an anchor before the Unix epoch, is accepted. Its window `(-301s, -1s]` holds no +-- sample, so the result is empty. +TQL EVAL (0, 0, '1s') at_modifier_gauge @ -1; + ++----+-----+------+ +| ts | val | host | ++----+-----+------+ ++----+-----+------+ + +-- 8. `@` on a metric that does not exist yields an empty result. +TQL EVAL (100, 100, '1s') at_modifier_missing @ 300; + ++------+-------+ +| time | value | ++------+-------+ ++------+-------+ + +-- The same for a range selector over a missing metric: the anchored window is folded once by the +-- promoted call, but the scan has no row to fold, so the whole grid stays empty. +TQL EVAL (0, 240, '60s') rate(at_modifier_missing[5m] @ 300); + ++------+------------------------------------------------+ +| time | prom_rate(time_range,value,time,Int64(300000)) | ++------+------------------------------------------------+ ++------+------------------------------------------------+ + +-- 9. Multi-series coverage of the promoted call. The call above the anchored range selector is +-- evaluated once per series and its result is replayed per series, so grouping by `host` must +-- report both groups at every step: 1.0 for host 'a' and 0.5 for host 'b'. A promoted result that +-- treated the two series as one timeline would report the first step from the first row of the +-- batch and the later steps from its last row, dropping or mixing a host here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)); + ++------+---------------------+-----------------------------------------------+ +| host | ts | sum(prom_rate(ts_range,val,ts,Int64(300000))) | ++------+---------------------+-----------------------------------------------+ +| a | 1970-01-01T00:00:00 | 1.0 | +| a | 1970-01-01T00:01:00 | 1.0 | +| a | 1970-01-01T00:02:00 | 1.0 | +| a | 1970-01-01T00:03:00 | 1.0 | +| a | 1970-01-01T00:04:00 | 1.0 | +| b | 1970-01-01T00:00:00 | 0.5 | +| b | 1970-01-01T00:01:00 | 0.5 | +| b | 1970-01-01T00:02:00 | 0.5 | +| b | 1970-01-01T00:03:00 | 0.5 | +| b | 1970-01-01T00:04:00 | 0.5 | ++------+---------------------+-----------------------------------------------+ + +-- Consistency: comparing the per-host sequence above against 1.0 keeps exactly host 'b' at +-- every step, confirming it stays at 0.5 for the whole grid. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)) != 1.0; + ++------+---------------------+-----------------------------------------------+ +| host | ts | sum(prom_rate(ts_range,val,ts,Int64(300000))) | ++------+---------------------+-----------------------------------------------+ +| b | 1970-01-01T00:00:00 | 0.5 | +| b | 1970-01-01T00:01:00 | 0.5 | +| b | 1970-01-01T00:02:00 | 0.5 | +| b | 1970-01-01T00:03:00 | 0.5 | +| b | 1970-01-01T00:04:00 | 0.5 | ++------+---------------------+-----------------------------------------------+ + +-- A value function directly over an anchored instant selector is not promoted: the selector +-- anchors and replays its sample per series, and `abs` is evaluated at every step over that replay. +-- It must keep one row per series per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 300); + ++---------------------+----------+------+ +| ts | abs(val) | host | ++---------------------+----------+------+ +| 1970-01-01T00:01:40 | 12.0 | b | +| 1970-01-01T00:01:40 | 5.0 | a | +| 1970-01-01T00:03:20 | 12.0 | b | +| 1970-01-01T00:03:20 | 5.0 | a | +| 1970-01-01T00:05:00 | 12.0 | b | +| 1970-01-01T00:05:00 | 5.0 | a | +| 1970-01-01T00:06:40 | 12.0 | b | +| 1970-01-01T00:06:40 | 5.0 | a | ++---------------------+----------+------+ + +-- 10. `predict_linear` over an anchored range selector: the regression is centered on the +-- evaluation instant of the step, not on the last sample of the window. The window `(0s, 300s]` of +-- host 'a' is left-open (the sample at 0s is outside it) and rises by 1.0 per minute, so each step +-- predicts 60s ahead of the trend read at its own instant: 6.0 at 300s and 7.0 at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 360, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 300, 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:00 | 6.0 | a | +| 1970-01-01T00:06:00 | 7.0 | a | ++---------------------+-----------------------------------------------+------+ + +-- The anchor and the start of the evaluation differ here, so the window `(-60s, 240s]` — samples 0s +-- to 240s — is shifted by `at_offset` = 300s - 240s = 60s. Every step must follow the evaluation +-- grid (6.0 to 11.0); a regression read at the wrong instant predicts 60s ahead of the window's +-- last sample instead, i.e. 5.0 at every step — the sample at 240s is 4.0 and the trend rises by +-- 1.0 per minute — and shows up as a constant here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 600, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 240, 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:00 | 6.0 | a | +| 1970-01-01T00:06:00 | 7.0 | a | +| 1970-01-01T00:07:00 | 8.0 | a | +| 1970-01-01T00:08:00 | 9.0 | a | +| 1970-01-01T00:09:00 | 10.0 | a | +| 1970-01-01T00:10:00 | 11.0 | a | ++---------------------+-----------------------------------------------+------+ + +-- Control: without `@` every step folds its own window, whose last sample is up to a step older +-- than the evaluation instant. Reading the regression off that sample reports 6.0 / 7.0 instead +-- of 6.5 / 7.5 here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m], 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:30 | 6.5 | a | +| 1970-01-01T00:06:30 | 7.5 | a | ++---------------------+-----------------------------------------------+------+ + +-- The same holds with an `offset`, which shifts the whole window backwards without moving the +-- evaluation instants: the window of step `T` is the left-open `(T - 1m - 5m, T - 1m]`, and the +-- trend is read at `T`, i.e. one minute past its newest sample. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] offset 1m, 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:30 | 6.5 | a | +| 1970-01-01T00:06:30 | 7.5 | a | ++---------------------+-----------------------------------------------+------+ + +-- `timestamp()` of an anchored instant selector is not promoted either: its argument is an +-- instant selector, not a range selector, so the selector anchors and replays its sample per series +-- on its own and `timestamp()` reports the timestamp of the sample the anchor selected at every +-- step (300s here). Its plan keeps the anchored selection below the replay of the grid and projects +-- the timestamp above it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + ++---------------------+-------+------+ +| ts | value | host | ++---------------------+-------+------+ +| 1970-01-01T00:00:00 | 300.0 | a | +| 1970-01-01T00:00:00 | 300.0 | b | +| 1970-01-01T00:01:00 | 300.0 | a | +| 1970-01-01T00:01:00 | 300.0 | b | +| 1970-01-01T00:02:00 | 300.0 | a | +| 1970-01-01T00:02:00 | 300.0 | b | +| 1970-01-01T00:03:00 | 300.0 | a | +| 1970-01-01T00:03:00 | 300.0 | b | +| 1970-01-01T00:04:00 | 300.0 | a | +| 1970-01-01T00:04:00 | 300.0 | b | ++---------------------+-------+------+ + +-- With an `offset`, the anchor selects the sample at 240s, and `timestamp()` reports 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 120, '60s') timestamp(at_modifier_gauge @ 300 offset 1m); + ++---------------------+-------+------+ +| ts | value | host | ++---------------------+-------+------+ +| 1970-01-01T00:00:00 | 240.0 | a | +| 1970-01-01T00:00:00 | 240.0 | b | +| 1970-01-01T00:01:00 | 240.0 | a | +| 1970-01-01T00:01:00 | 240.0 | b | +| 1970-01-01T00:02:00 | 240.0 | a | +| 1970-01-01T00:02:00 | 240.0 | b | ++---------------------+-------+------+ + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +TQL EXPLAIN (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: at_modifier_gauge.ts, value AS value, at_modifier_gauge.host | +| | Projection: at_modifier_gauge.ts, __promql_timestamp_value_ AS value, at_modifier_gauge.host | +| | Filter: at_modifier_gauge.val IS NOT NULL | +| | PromInstantManipulate: range=[0..240000], lookback=[240001], interval=[60000], time index=[ts] | +| | PromInstantManipulate: range=[0..0], lookback=[300000], interval=[60000], time index=[ts] | +| | Projection: at_modifier_gauge.ts, at_modifier_gauge.val, at_modifier_gauge.host, CAST(CAST(at_modifier_gauge.ts AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ | +| | PromSeriesNormalize: offset=[-300000], time index=[ts], filter NaN: [false] | +| | PromSeriesDivide: tags=["host"] | +| | Sort: at_modifier_gauge.host ASC NULLS FIRST, at_modifier_gauge.ts ASC NULLS FIRST | +| | Filter: at_modifier_gauge.ts >= TimestampMillisecond(1, None) AND at_modifier_gauge.ts <= TimestampMillisecond(300000, None) | +| | TableScan: at_modifier_gauge, partial_filters=[at_modifier_gauge.ts >= TimestampMillisecond(1, None), at_modifier_gauge.ts <= TimestampMillisecond(300000, None)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- 11. `label_join` above an anchored selector. The call rewrites the label the input series are +-- told apart by (`host` becomes `a-a` and `b-b`), so it is never promoted: it must stay above the +-- per-series replay of the selector and be evaluated at every step. Both hosts are still reported, +-- with their own value, at every step. A promoted `label_join` would instead be replayed through +-- the labels it just rewrote and report the two series as a single timeline, i.e. one mixed row +-- per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "-", "host", "host"); + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:05:00 | 12.0 | b-b | +| 1970-01-01T00:05:00 | 5.0 | a-a | +| 1970-01-01T00:06:00 | 12.0 | b-b | +| 1970-01-01T00:06:00 | 5.0 | a-a | +| 1970-01-01T00:07:00 | 12.0 | b-b | +| 1970-01-01T00:07:00 | 5.0 | a-a | +| 1970-01-01T00:08:00 | 12.0 | b-b | +| 1970-01-01T00:08:00 | 5.0 | a-a | ++---------------------+------+------+ + +-- The same join below another call: neither is promoted, and both are evaluated at every step over +-- the per-series replay of the anchored selector. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "-", "host", "host")); + ++---------------------+----------+------+ +| ts | abs(val) | host | ++---------------------+----------+------+ +| 1970-01-01T00:05:00 | 12.0 | b-b | +| 1970-01-01T00:05:00 | 5.0 | a-a | +| 1970-01-01T00:06:00 | 12.0 | b-b | +| 1970-01-01T00:06:00 | 5.0 | a-a | +| 1970-01-01T00:07:00 | 12.0 | b-b | +| 1970-01-01T00:07:00 | 5.0 | a-a | +| 1970-01-01T00:08:00 | 12.0 | b-b | +| 1970-01-01T00:08:00 | 5.0 | a-a | ++---------------------+----------+------+ + +-- A range call below the join is still promoted on its own (it is the direct call over the +-- anchored range selector): the anchored window is folded once per series, and the join above that +-- replay reports the rate of both hosts at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "-", "host", "host"); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:05:00 | 0.5 | b-b | +| 1970-01-01T00:05:00 | 1.0 | a-a | +| 1970-01-01T00:06:00 | 0.5 | b-b | +| 1970-01-01T00:06:00 | 1.0 | a-a | +| 1970-01-01T00:07:00 | 0.5 | b-b | +| 1970-01-01T00:07:00 | 1.0 | a-a | +| 1970-01-01T00:08:00 | 0.5 | b-b | +| 1970-01-01T00:08:00 | 1.0 | a-a | ++---------------------+------------------------------------------+------+ + +DROP TABLE at_modifier_gauge; + +Affected Rows: 0 + +DROP TABLE at_modifier_counter_total; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/at_modifier.sql b/tests/cases/standalone/common/promql/at_modifier.sql new file mode 100644 index 00000000000..b4fd9d1a552 --- /dev/null +++ b/tests/cases/standalone/common/promql/at_modifier.sql @@ -0,0 +1,296 @@ +-- Tests for the PromQL `@` modifier on vector and matrix selectors. +-- +-- `@` anchors the sample selection window at a fixed timestamp instead of the timestamp of each +-- evaluation step: `@ ` uses the given timestamp, `@ start()` / `@ end()` use the +-- start/end of the statement's evaluation range, and `offset` shifts the anchor backwards before +-- the window is selected. The output timestamps still follow the evaluation grid. +-- +-- Every metric below carries two series, host 'a' and host 'b', with different values: the +-- anchored results must therefore report both of them, per step. A series silently dropped, or a +-- batch holding several series handled as one timeline, does not look like a correct single-series +-- answer here but shows up as a missing or shifted row. +-- +-- Sample timestamps below are in milliseconds, while `@` and `TQL EVAL` timestamps are in seconds, +-- as in Prometheus. + +CREATE TABLE at_modifier_gauge ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +CREATE TABLE at_modifier_counter_total ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +-- One sample per minute, so the anchored sample is easy to identify. The default lookback is 5m. +-- Host 'a' reports 0.0 to 6.0, host 'b' reports 7.0 to 13.0: the two series are told apart by their +-- values alone, which stays true after any anchoring. +INSERT INTO at_modifier_gauge VALUES + (0, 0.0, 'a'), + (60000, 1.0, 'a'), + (120000, 2.0, 'a'), + (180000, 3.0, 'a'), + (240000, 4.0, 'a'), + (300000, 5.0, 'a'), + (360000, 6.0, 'a'), + (0, 7.0, 'b'), + (60000, 8.0, 'b'), + (120000, 9.0, 'b'), + (180000, 10.0, 'b'), + (240000, 11.0, 'b'), + (300000, 12.0, 'b'), + (360000, 13.0, 'b'); + +-- Host 'a': a counter increasing by 60 per minute, i.e. by 1 per second. +-- Host 'b': a counter starting at 100 and increasing by 30 per minute, i.e. by 0.5 per second. Its +-- rate differs from host 'a', so a window folded for the wrong series cannot pass as the right one. +INSERT INTO at_modifier_counter_total VALUES + (0, 0.0, 'a'), + (60000, 60.0, 'a'), + (120000, 120.0, 'a'), + (180000, 180.0, 'a'), + (240000, 240.0, 'a'), + (300000, 300.0, 'a'), + (360000, 360.0, 'a'), + (0, 100.0, 'b'), + (60000, 130.0, 'b'), + (120000, 160.0, 'b'), + (180000, 190.0, 'b'), + (240000, 220.0, 'b'), + (300000, 250.0, 'b'), + (360000, 280.0, 'b'); + +-- 1. Instant query anchored at an absolute timestamp: the sample window is `(anchor - lookback, +-- anchor]` = `(0s, 300s]`, so its newest samples are 5.0 for host 'a' and 12.0 for host 'b', both +-- at 300s. The output timestamps stay the evaluation timestamps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300; + +-- Control: the same query without `@` looks back from the evaluation timestamp instead: +-- `(-200s, 100s]` selects the samples 1.0 (a) and 8.0 (b) at 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge; + +-- The anchor, not the evaluation timestamp, decides which sample is reported. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 360; + +-- Sub-second precision is preserved in the anchor: `240.999` selects the samples at 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 240.999; + +-- The window `(-300s, 0s]` includes the samples exactly at the anchor. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 0; + +-- 2. `@ start()`: every step selects its samples around the start of the evaluation range +-- (`(-200s, 100s]`), so every step reports the same value per series (1.0 for 'a', 8.0 for 'b'). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ start(); + +-- Control: without `@` every step looks back from its own timestamp, so the steps differ. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge; + +-- 3. `@ end()`: every step selects its samples around the end of the evaluation range +-- (`(100s, 400s]`), whose newest samples are 6.0 (a) and 13.0 (b) at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ end(); + +-- 4. Range selector anchored by `@`: the call directly above the anchored range selector is +-- evaluated once, at the start of the evaluation, and its result is reported at every step — the +-- planner's counterpart of Prometheus' `StepInvariantExpr`. The window `(0s, 300s]` is left-open, +-- so it holds 60 -> 300 for host 'a' and 130 -> 250 for host 'b', i.e. 240 and 120 units. The 60s +-- gap to the window start is below the extrapolation threshold (1.1 x the 60s sampling interval), +-- so `rate` adds it and divides by the 300s window: 240 * 1.25 / 300 = 1.0 and +-- 120 * 1.25 / 300 = 0.5 per second. Every step must report those same per-series values, and both +-- series must be there at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300); + +-- Consistency: the instant query below shows the anchored values over the same left-open +-- `(0s, 300s]` window (1.0 per second for 'a', 0.5 for 'b'). The `!= 1.0` comparison filters +-- host 'a', but host 'b' remains at every step with its rate of 0.5 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 300, '1s') rate(at_modifier_counter_total[5m]); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300) != 1.0; + +-- Parentheses around the range argument are transparent: the call must be promoted exactly like +-- the unparenthesized one above, so every step reports the same anchored values for both series +-- (1.0 for host 'a', 0.5 for host 'b') instead of folding a window per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate((at_modifier_counter_total[5m] @ 300)); + +-- The same holds for `@ start()` and `@ end()`, which are also fixed anchors: they resolve to the +-- statement's evaluation range, so a range call using them is evaluated once as well. +-- The statement starts at 60s: the window `(-240s, 60s]` holds only the samples at 0s and 60s, so +-- host 'a' goes 0 -> 60 and host 'b' 100 -> 130. The far 240s gap to the window start is beyond the +-- extrapolation threshold (1.1 x the 60s interval), so only half an interval (30s) is added there; +-- host 'a' starts its counter at 0, which clamps that back: 60 / 300 = 0.2 and +-- 30 * 1.5 / 300 = 0.15 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 300, '60s') rate(at_modifier_counter_total[5m] @ start()); + +-- The instant query at 60s folds the same window as every step of the query above. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 60, '1s') rate(at_modifier_counter_total[5m]); + +-- The statement ends at 240s: the window `(-60s, 240s]` holds the samples 0s to 240s, so host 'a' +-- goes 0 -> 240 and host 'b' 100 -> 220. The 60s gap to the window start is added for both, but +-- host 'a' starts at zero, which clamps it back (240 / 300 = 0.8 per second) while host 'b' keeps +-- it (120 * 1.25 / 300 = 0.5 per second). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ end()); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (240, 240, '1s') rate(at_modifier_counter_total[5m]); + +-- Control: without `@` each step folds its own window, so the values keep changing. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m]); + +-- `@ start()` on a range selector: the window `(-300s, 0s]` is left-open, so it holds a single +-- sample per series — the one exactly at the anchor — and `count_over_time` returns 1 at every step +-- for both of them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m] @ start()); + +-- Control: without `@` the window keeps moving, so the count grows with the step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m]); + +-- 5. `@` combined with `offset`: the offset moves the anchor backwards before the window is +-- selected, so `@ 300 offset 2m` reports the samples at 180s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 2m; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 1m; + +-- Control: the offset applies to the evaluation timestamp when there is no `@`. +TQL EVAL (100, 100, '1s') at_modifier_gauge offset 2m; + +-- 6. `@` inside a larger expression. The anchored operand contributes the same value at every +-- step, while the plain operand still follows the evaluation step. A binary expression is never +-- promoted: when one side is anchored and the other is not, the anchored side replays on its own, +-- and when both sides are anchored, each side reports its own anchored samples. The join then runs +-- at every step over the replayed per-series inputs. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge @ 0; + +-- The aggregation above the anchored selector runs at every step over the replayed (still +-- per-series) samples, so it sums both series at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') sum(at_modifier_gauge @ 0); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 0); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge; + +-- 7. `@ -1`, an anchor before the Unix epoch, is accepted. Its window `(-301s, -1s]` holds no +-- sample, so the result is empty. +TQL EVAL (0, 0, '1s') at_modifier_gauge @ -1; + +-- 8. `@` on a metric that does not exist yields an empty result. +TQL EVAL (100, 100, '1s') at_modifier_missing @ 300; + +-- The same for a range selector over a missing metric: the anchored window is folded once by the +-- promoted call, but the scan has no row to fold, so the whole grid stays empty. +TQL EVAL (0, 240, '60s') rate(at_modifier_missing[5m] @ 300); + +-- 9. Multi-series coverage of the promoted call. The call above the anchored range selector is +-- evaluated once per series and its result is replayed per series, so grouping by `host` must +-- report both groups at every step: 1.0 for host 'a' and 0.5 for host 'b'. A promoted result that +-- treated the two series as one timeline would report the first step from the first row of the +-- batch and the later steps from its last row, dropping or mixing a host here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)); + +-- Consistency: comparing the per-host sequence above against 1.0 keeps exactly host 'b' at +-- every step, confirming it stays at 0.5 for the whole grid. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)) != 1.0; + +-- A value function directly over an anchored instant selector is not promoted: the selector +-- anchors and replays its sample per series, and `abs` is evaluated at every step over that replay. +-- It must keep one row per series per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 300); + +-- 10. `predict_linear` over an anchored range selector: the regression is centered on the +-- evaluation instant of the step, not on the last sample of the window. The window `(0s, 300s]` of +-- host 'a' is left-open (the sample at 0s is outside it) and rises by 1.0 per minute, so each step +-- predicts 60s ahead of the trend read at its own instant: 6.0 at 300s and 7.0 at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 360, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 300, 60); + +-- The anchor and the start of the evaluation differ here, so the window `(-60s, 240s]` — samples 0s +-- to 240s — is shifted by `at_offset` = 300s - 240s = 60s. Every step must follow the evaluation +-- grid (6.0 to 11.0); a regression read at the wrong instant predicts 60s ahead of the window's +-- last sample instead, i.e. 5.0 at every step — the sample at 240s is 4.0 and the trend rises by +-- 1.0 per minute — and shows up as a constant here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 600, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 240, 60); + +-- Control: without `@` every step folds its own window, whose last sample is up to a step older +-- than the evaluation instant. Reading the regression off that sample reports 6.0 / 7.0 instead +-- of 6.5 / 7.5 here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m], 60); + +-- The same holds with an `offset`, which shifts the whole window backwards without moving the +-- evaluation instants: the window of step `T` is the left-open `(T - 1m - 5m, T - 1m]`, and the +-- trend is read at `T`, i.e. one minute past its newest sample. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] offset 1m, 60); + +-- `timestamp()` of an anchored instant selector is not promoted either: its argument is an +-- instant selector, not a range selector, so the selector anchors and replays its sample per series +-- on its own and `timestamp()` reports the timestamp of the sample the anchor selected at every +-- step (300s here). Its plan keeps the anchored selection below the replay of the grid and projects +-- the timestamp above it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + +-- With an `offset`, the anchor selects the sample at 240s, and `timestamp()` reports 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 120, '60s') timestamp(at_modifier_gauge @ 300 offset 1m); + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +TQL EXPLAIN (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + +-- 11. `label_join` above an anchored selector. The call rewrites the label the input series are +-- told apart by (`host` becomes `a-a` and `b-b`), so it is never promoted: it must stay above the +-- per-series replay of the selector and be evaluated at every step. Both hosts are still reported, +-- with their own value, at every step. A promoted `label_join` would instead be replayed through +-- the labels it just rewrote and report the two series as a single timeline, i.e. one mixed row +-- per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "-", "host", "host"); + +-- The same join below another call: neither is promoted, and both are evaluated at every step over +-- the per-series replay of the anchored selector. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "-", "host", "host")); + +-- A range call below the join is still promoted on its own (it is the direct call over the +-- anchored range selector): the anchored window is folded once per series, and the join above that +-- replay reports the rate of both hosts at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "-", "host", "host"); + +DROP TABLE at_modifier_gauge; + +DROP TABLE at_modifier_counter_total; diff --git a/tests/cases/standalone/common/promql/computed_scalar_rhs.result b/tests/cases/standalone/common/promql/computed_scalar_rhs.result new file mode 100644 index 00000000000..96b925d1cf5 --- /dev/null +++ b/tests/cases/standalone/common/promql/computed_scalar_rhs.result @@ -0,0 +1,97 @@ +-- A vector combined with a computed scalar on the right keeps the vector's +-- labels, like a scalar on the left does. Expected results are from Prometheus 3. +create table computed_scalar_lhs ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +create table computed_scalar_rhs ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +insert into computed_scalar_lhs values + (0, 1.0, 'a'), + (0, 2.0, 'b'); + +Affected Rows: 2 + +insert into computed_scalar_rhs values + (0, 10.0, 'x'); + +Affected Rows: 1 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs + time(); + ++------+---------------------+-------------------------------------------------+ +| host | ts | computed_scalar_lhs.val + .time / Float64(1000) | ++------+---------------------+-------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 61.0 | +| b | 1970-01-01T00:01:00 | 62.0 | ++------+---------------------+-------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') time() + computed_scalar_lhs; + ++------+---------------------+-------------------------------------------------+ +| host | ts | .time / Float64(1000) + computed_scalar_lhs.val | ++------+---------------------+-------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 61.0 | +| b | 1970-01-01T00:01:00 | 62.0 | ++------+---------------------+-------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs); + ++------+---------------------+-----------------------------------------------------------+ +| host | ts | computed_scalar_lhs.val * computed_scalar_rhs.scalar(val) | ++------+---------------------+-----------------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 10.0 | +| b | 1970-01-01T00:01:00 | 20.0 | ++------+---------------------+-----------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2); + ++------+---------------------+--------------------------------------------------------------+ +| host | ts | computed_scalar_lhs.val < .time / Float64(1000) * Float64(2) | ++------+---------------------+--------------------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 1.0 | +| b | 1970-01-01T00:01:00 | 1.0 | ++------+---------------------+--------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs); + ++------+---------------------+---------------------------------------------+ +| host | ts | power(.val,computed_scalar_rhs.scalar(val)) | ++------+---------------------+---------------------------------------------+ +| a | 1970-01-01T00:01:00 | 1.0 | +| b | 1970-01-01T00:01:00 | 1024.0 | ++------+---------------------+---------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs); + ++------+---------------------+----------------------------------------+ +| host | ts | computed_scalar_rhs.scalar(val) * .val | ++------+---------------------+----------------------------------------+ +| a | 1970-01-01T00:01:00 | 10.0 | +| b | 1970-01-01T00:01:00 | 20.0 | ++------+---------------------+----------------------------------------+ + +drop table computed_scalar_lhs; + +Affected Rows: 0 + +drop table computed_scalar_rhs; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/computed_scalar_rhs.sql b/tests/cases/standalone/common/promql/computed_scalar_rhs.sql new file mode 100644 index 00000000000..57cc431c1f9 --- /dev/null +++ b/tests/cases/standalone/common/promql/computed_scalar_rhs.sql @@ -0,0 +1,43 @@ +-- A vector combined with a computed scalar on the right keeps the vector's +-- labels, like a scalar on the left does. Expected results are from Prometheus 3. + +create table computed_scalar_lhs ( + ts timestamp time index, + val double, + host string primary key +); + +create table computed_scalar_rhs ( + ts timestamp time index, + val double, + host string primary key +); + +insert into computed_scalar_lhs values + (0, 1.0, 'a'), + (0, 2.0, 'b'); + +insert into computed_scalar_rhs values + (0, 10.0, 'x'); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs + time(); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') time() + computed_scalar_lhs; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs); + +drop table computed_scalar_lhs; + +drop table computed_scalar_rhs; diff --git a/tests/cases/standalone/common/promql/count_values.result b/tests/cases/standalone/common/promql/count_values.result index e2f65eae342..886c3a143b7 100644 --- a/tests/cases/standalone/common/promql/count_values.result +++ b/tests/cases/standalone/common/promql/count_values.result @@ -61,6 +61,145 @@ TQL EVAL (0, 15, '5s') count_values("status_code", http_requests) by (idc); | 2 | idc2 | 1970-01-01T00:00:15 | 500 | +--------------------------+------+---------------------+-------------+ +-- https://github.com/GreptimeTeam/greptimedb/issues/9181 +-- The label generated by `count_values` must survive an enclosing expression, +-- e.g. a per-sample function, an arithmetic operation or an enclosing aggregation. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') abs(count_values("v", http_requests)); + ++---------------------+-------------------------------+-----+ +| ts | abs(count(http_requests.val)) | v | ++---------------------+-------------------------------+-----+ +| 1970-01-01T00:00:00 | 1 | 401 | +| 1970-01-01T00:00:00 | 3 | 200 | +| 1970-01-01T00:00:05 | 1 | 401 | +| 1970-01-01T00:00:05 | 1 | 500 | +| 1970-01-01T00:00:05 | 2 | 404 | +| 1970-01-01T00:00:10 | 2 | 200 | +| 1970-01-01T00:00:10 | 2 | 201 | +| 1970-01-01T00:00:15 | 4 | 500 | ++---------------------+-------------------------------+-----+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') round(count_values("v", http_requests)); + ++---------------------+-------------------------------------------------+-----+ +| ts | prom_round(count(http_requests.val),Float64(0)) | v | ++---------------------+-------------------------------------------------+-----+ +| 1970-01-01T00:00:00 | 1.0 | 401 | +| 1970-01-01T00:00:00 | 3.0 | 200 | +| 1970-01-01T00:00:05 | 1.0 | 401 | +| 1970-01-01T00:00:05 | 1.0 | 500 | +| 1970-01-01T00:00:05 | 2.0 | 404 | +| 1970-01-01T00:00:10 | 2.0 | 200 | +| 1970-01-01T00:00:10 | 2.0 | 201 | +| 1970-01-01T00:00:15 | 4.0 | 500 | ++---------------------+-------------------------------------------------+-----+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("v", http_requests) + 1; + ++-----+---------------------+---------------------------------------+ +| v | ts | count(http_requests.val) + Float64(1) | ++-----+---------------------+---------------------------------------+ +| 200 | 1970-01-01T00:00:00 | 4.0 | +| 200 | 1970-01-01T00:00:10 | 3.0 | +| 201 | 1970-01-01T00:00:10 | 3.0 | +| 401 | 1970-01-01T00:00:00 | 2.0 | +| 401 | 1970-01-01T00:00:05 | 2.0 | +| 404 | 1970-01-01T00:00:05 | 3.0 | +| 500 | 1970-01-01T00:00:05 | 2.0 | +| 500 | 1970-01-01T00:00:15 | 5.0 | ++-----+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') sum by (v) (count_values("v", http_requests)); + ++-----+---------------------+-------------------------------+ +| v | ts | sum(count(http_requests.val)) | ++-----+---------------------+-------------------------------+ +| 200 | 1970-01-01T00:00:00 | 3 | +| 200 | 1970-01-01T00:00:10 | 2 | +| 201 | 1970-01-01T00:00:10 | 2 | +| 401 | 1970-01-01T00:00:00 | 1 | +| 401 | 1970-01-01T00:00:05 | 1 | +| 404 | 1970-01-01T00:00:05 | 2 | +| 500 | 1970-01-01T00:00:05 | 1 | +| 500 | 1970-01-01T00:00:15 | 4 | ++-----+---------------------+-------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') topk(10, count_values("v", http_requests)); + ++--------------------------+-----+---------------------+ +| count(http_requests.val) | v | ts | ++--------------------------+-----+---------------------+ +| 1 | 401 | 1970-01-01T00:00:00 | +| 1 | 401 | 1970-01-01T00:00:05 | +| 1 | 500 | 1970-01-01T00:00:05 | +| 2 | 200 | 1970-01-01T00:00:10 | +| 2 | 201 | 1970-01-01T00:00:10 | +| 2 | 404 | 1970-01-01T00:00:05 | +| 3 | 200 | 1970-01-01T00:00:00 | +| 4 | 500 | 1970-01-01T00:00:15 | ++--------------------------+-----+---------------------+ + +-- The generated label overwrites an input label with the same name (PromQL semantics), +-- and must not turn the generated column into an ambiguous reference. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("idc", http_requests) by (idc); + ++--------------------------+---------------------+-----+ +| count(http_requests.val) | ts | idc | ++--------------------------+---------------------+-----+ +| 1 | 1970-01-01T00:00:00 | 401 | +| 1 | 1970-01-01T00:00:05 | 401 | +| 1 | 1970-01-01T00:00:05 | 500 | +| 2 | 1970-01-01T00:00:05 | 404 | +| 2 | 1970-01-01T00:00:10 | 200 | +| 2 | 1970-01-01T00:00:10 | 201 | +| 3 | 1970-01-01T00:00:00 | 200 | +| 4 | 1970-01-01T00:00:15 | 500 | ++--------------------------+---------------------+-----+ + +-- The generated label is the sample value in PromQL's textual form +-- (`strconv.FormatFloat(value, 'f', -1, 64)`: the shortest decimal form without an +-- exponent), so it is a string: `200.0` becomes the label `200` and `1e21` becomes +-- `1000000000000000000000`. +CREATE TABLE http_requests_double ( + ts timestamp(3) time index, + host STRING, + idc STRING, + val DOUBLE, + PRIMARY KEY(host, idc), +); + +Affected Rows: 0 + +INSERT INTO TABLE http_requests_double VALUES + (0, 'host1', "idc1", 200.0), + (0, 'host2', "idc1", 0.5), + (5000, 'host1', "idc1", 200.0), + (5000, 'host2', "idc1", 1e21); + +Affected Rows: 4 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') count_values("val_label", http_requests_double); + ++---------------------------------+---------------------+------------------------+ +| count(http_requests_double.val) | ts | val_label | ++---------------------------------+---------------------+------------------------+ +| 1 | 1970-01-01T00:00:00 | 0.5 | +| 1 | 1970-01-01T00:00:00 | 200 | +| 1 | 1970-01-01T00:00:05 | 1000000000000000000000 | +| 1 | 1970-01-01T00:00:05 | 200 | ++---------------------------------+---------------------+------------------------+ + +DROP TABLE http_requests_double; + +Affected Rows: 0 + DROP TABLE http_requests; Affected Rows: 0 diff --git a/tests/cases/standalone/common/promql/count_values.sql b/tests/cases/standalone/common/promql/count_values.sql index 09c2b136b83..f23cdec5fbe 100644 --- a/tests/cases/standalone/common/promql/count_values.sql +++ b/tests/cases/standalone/common/promql/count_values.sql @@ -28,4 +28,50 @@ TQL EVAL (0, 15, '5s') count_values("status_code", http_requests); TQL EVAL (0, 15, '5s') count_values("status_code", http_requests) by (idc); +-- https://github.com/GreptimeTeam/greptimedb/issues/9181 +-- The label generated by `count_values` must survive an enclosing expression, +-- e.g. a per-sample function, an arithmetic operation or an enclosing aggregation. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') abs(count_values("v", http_requests)); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') round(count_values("v", http_requests)); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("v", http_requests) + 1; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') sum by (v) (count_values("v", http_requests)); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') topk(10, count_values("v", http_requests)); + +-- The generated label overwrites an input label with the same name (PromQL semantics), +-- and must not turn the generated column into an ambiguous reference. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("idc", http_requests) by (idc); + +-- The generated label is the sample value in PromQL's textual form +-- (`strconv.FormatFloat(value, 'f', -1, 64)`: the shortest decimal form without an +-- exponent), so it is a string: `200.0` becomes the label `200` and `1e21` becomes +-- `1000000000000000000000`. +CREATE TABLE http_requests_double ( + ts timestamp(3) time index, + host STRING, + idc STRING, + val DOUBLE, + PRIMARY KEY(host, idc), +); + +INSERT INTO TABLE http_requests_double VALUES + (0, 'host1', "idc1", 200.0), + (0, 'host2', "idc1", 0.5), + (5000, 'host1', "idc1", 200.0), + (5000, 'host2', "idc1", 1e21); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') count_values("val_label", http_requests_double); + +DROP TABLE http_requests_double; + DROP TABLE http_requests; diff --git a/tests/cases/standalone/common/promql/dotted_column_names.result b/tests/cases/standalone/common/promql/dotted_column_names.result new file mode 100644 index 00000000000..c855a5ba5c7 --- /dev/null +++ b/tests/cases/standalone/common/promql/dotted_column_names.result @@ -0,0 +1,505 @@ +-- https://github.com/GreptimeTeam/greptimedb/issues/9390 +-- Tag, field and time index names containing dots or upper case letters must +-- be resolved as plain column names, not as `relation.column`. `Host` and `host` +-- hold different values, so reading one for the other changes the output. +CREATE TABLE "otel.m" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + "Host" STRING, + host STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", "Host", host) +) PARTITION ON COLUMNS ("service.name") ( + "service.name" < 'b', + "service.name" >= 'b' +); + +Affected Rows: 0 + +INSERT INTO "otel.m" VALUES + (0, 'a', 'h1', 'x', 1), + (0, 'a', 'h2', 'x', 2), + (0, 'b', 'h1', 'y', 3), + (5000, 'a', 'h1', 'x', 4), + (5000, 'a', 'h2', 'x', 5), + (5000, 'b', 'h1', 'y', 6), + (10000, 'a', 'h1', 'x', 8), + (10000, 'a', 'h2', 'x', 9), + (10000, 'b', 'h1', 'y', 8); + +Affected Rows: 9 + +CREATE TABLE "otel.h" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + le STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", le) +); + +Affected Rows: 0 + +INSERT INTO "otel.h" VALUES + (0, 'a', '0.1', 1), + (0, 'a', '1', 3), + (0, 'a', '+Inf', 4), + (10000, 'a', '0.1', 2), + (10000, 'a', '1', 6), + (10000, 'a', '+Inf', 10); + +Affected Rows: 6 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | a | h1 | x | 1.0 | +| 1970-01-01T00:00:00 | b | h1 | y | 3.0 | +| 1970-01-01T00:00:05 | a | h1 | x | 4.0 | +| 1970-01-01T00:00:05 | b | h1 | y | 6.0 | +| 1970-01-01T00:00:10 | a | h1 | x | 8.0 | +| 1970-01-01T00:00:10 | b | h1 | y | 8.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on("service.name") {"otel.m", "service.name"="b"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | b | h1 | y | 3.0 | +| 1970-01-01T00:00:05 | b | h1 | y | 6.0 | +| 1970-01-01T00:00:10 | b | h1 | y | 8.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on(host) {"otel.m", host="y"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | b | h1 | y | 3.0 | +| 1970-01-01T00:00:05 | b | h1 | y | 6.0 | +| 1970-01-01T00:00:10 | b | h1 | y | 8.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless {"otel.m", "Host"="h1"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | a | h2 | x | 2.0 | +| 1970-01-01T00:00:05 | a | h2 | x | 5.0 | +| 1970-01-01T00:00:10 | a | h2 | x | 9.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless on("service.name") {"otel.m", "service.name"="b"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | a | h1 | x | 1.0 | +| 1970-01-01T00:00:00 | a | h2 | x | 2.0 | +| 1970-01-01T00:00:05 | a | h1 | x | 4.0 | +| 1970-01-01T00:00:05 | a | h2 | x | 5.0 | +| 1970-01-01T00:00:10 | a | h1 | x | 8.0 | +| 1970-01-01T00:00:10 | a | h2 | x | 9.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} or {"otel.m", "Host"="h1"}; + ++---------------------+------+------+--------------+-------+ +| ts.time | Host | host | service.name | v.val | ++---------------------+------+------+--------------+-------+ +| 1970-01-01T00:00:00 | h1 | x | a | 1.0 | +| 1970-01-01T00:00:00 | h1 | y | b | 3.0 | +| 1970-01-01T00:00:00 | h2 | x | a | 2.0 | +| 1970-01-01T00:00:05 | h1 | x | a | 4.0 | +| 1970-01-01T00:00:05 | h1 | y | b | 6.0 | +| 1970-01-01T00:00:05 | h2 | x | a | 5.0 | +| 1970-01-01T00:00:10 | h1 | x | a | 8.0 | +| 1970-01-01T00:00:10 | h1 | y | b | 8.0 | +| 1970-01-01T00:00:10 | h2 | x | a | 9.0 | ++---------------------+------+------+--------------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values("v.name", {"otel.m"}); + ++---------------------+---------------------+--------+ +| count(otel.m.v.val) | ts.time | v.name | ++---------------------+---------------------+--------+ +| 1 | 1970-01-01T00:00:00 | 1 | +| 1 | 1970-01-01T00:00:00 | 2 | +| 1 | 1970-01-01T00:00:00 | 3 | +| 1 | 1970-01-01T00:00:05 | 4 | +| 1 | 1970-01-01T00:00:05 | 5 | +| 1 | 1970-01-01T00:00:05 | 6 | +| 1 | 1970-01-01T00:00:10 | 9 | +| 2 | 1970-01-01T00:00:10 | 8 | ++---------------------+---------------------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values by ("service.name") ("v.name", {"otel.m"}); + ++---------------------+--------------+---------------------+--------+ +| count(otel.m.v.val) | service.name | ts.time | v.name | ++---------------------+--------------+---------------------+--------+ +| 1 | a | 1970-01-01T00:00:00 | 1 | +| 1 | a | 1970-01-01T00:00:00 | 2 | +| 1 | a | 1970-01-01T00:00:05 | 4 | +| 1 | a | 1970-01-01T00:00:05 | 5 | +| 1 | a | 1970-01-01T00:00:10 | 8 | +| 1 | a | 1970-01-01T00:00:10 | 9 | +| 1 | b | 1970-01-01T00:00:00 | 3 | +| 1 | b | 1970-01-01T00:00:05 | 6 | +| 1 | b | 1970-01-01T00:00:10 | 8 | ++---------------------+--------------+---------------------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') topk(1, {"otel.m"}); + ++-------+--------------+------+------+---------------------+ +| v.val | service.name | Host | host | ts.time | ++-------+--------------+------+------+---------------------+ +| 3.0 | b | h1 | y | 1970-01-01T00:00:00 | +| 6.0 | b | h1 | y | 1970-01-01T00:00:05 | +| 9.0 | a | h2 | x | 1970-01-01T00:00:10 | ++-------+--------------+------+------+---------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') bottomk by ("service.name") (1, {"otel.m"}); + ++-------+--------------+------+------+---------------------+ +| v.val | service.name | Host | host | ts.time | ++-------+--------------+------+------+---------------------+ +| 1.0 | a | h1 | x | 1970-01-01T00:00:00 | +| 3.0 | b | h1 | y | 1970-01-01T00:00:00 | +| 4.0 | a | h1 | x | 1970-01-01T00:00:05 | +| 6.0 | b | h1 | y | 1970-01-01T00:00:05 | +| 8.0 | a | h1 | x | 1970-01-01T00:00:10 | +| 8.0 | b | h1 | y | 1970-01-01T00:00:10 | ++-------+--------------+------+------+---------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("service.name") ({"otel.m"}); + ++--------------+---------------------+-------------------+ +| service.name | ts.time | sum(otel.m.v.val) | ++--------------+---------------------+-------------------+ +| a | 1970-01-01T00:00:00 | 3.0 | +| a | 1970-01-01T00:00:05 | 9.0 | +| a | 1970-01-01T00:00:10 | 17.0 | +| b | 1970-01-01T00:00:00 | 3.0 | +| b | 1970-01-01T00:00:05 | 6.0 | +| b | 1970-01-01T00:00:10 | 8.0 | ++--------------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("Host") ({"otel.m"}); + ++------+---------------------+-------------------+ +| Host | ts.time | sum(otel.m.v.val) | ++------+---------------------+-------------------+ +| h1 | 1970-01-01T00:00:00 | 4.0 | +| h1 | 1970-01-01T00:00:05 | 10.0 | +| h1 | 1970-01-01T00:00:10 | 16.0 | +| h2 | 1970-01-01T00:00:00 | 2.0 | +| h2 | 1970-01-01T00:00:05 | 5.0 | +| h2 | 1970-01-01T00:00:10 | 9.0 | ++------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by (host) ({"otel.m"}); + ++------+---------------------+-------------------+ +| host | ts.time | sum(otel.m.v.val) | ++------+---------------------+-------------------+ +| x | 1970-01-01T00:00:00 | 3.0 | +| x | 1970-01-01T00:00:05 | 9.0 | +| x | 1970-01-01T00:00:10 | 17.0 | +| y | 1970-01-01T00:00:00 | 3.0 | +| y | 1970-01-01T00:00:05 | 6.0 | +| y | 1970-01-01T00:00:10 | 8.0 | ++------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum without ("Host") ({"otel.m"}); + ++------+--------------+---------------------+-------------------+ +| host | service.name | ts.time | sum(otel.m.v.val) | ++------+--------------+---------------------+-------------------+ +| x | a | 1970-01-01T00:00:00 | 3.0 | +| x | a | 1970-01-01T00:00:05 | 9.0 | +| x | a | 1970-01-01T00:00:10 | 17.0 | +| y | b | 1970-01-01T00:00:00 | 3.0 | +| y | b | 1970-01-01T00:00:05 | 6.0 | +| y | b | 1970-01-01T00:00:10 | 8.0 | ++------+--------------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') -{"otel.m"}; + ++--------------+------+------+---------------------+-----------+ +| service.name | Host | host | ts.time | (- v.val) | ++--------------+------+------+---------------------+-----------+ +| a | h1 | x | 1970-01-01T00:00:00 | -1.0 | +| a | h1 | x | 1970-01-01T00:00:05 | -4.0 | +| a | h1 | x | 1970-01-01T00:00:10 | -8.0 | +| a | h2 | x | 1970-01-01T00:00:00 | -2.0 | +| a | h2 | x | 1970-01-01T00:00:05 | -5.0 | +| a | h2 | x | 1970-01-01T00:00:10 | -9.0 | +| b | h1 | y | 1970-01-01T00:00:00 | -3.0 | +| b | h1 | y | 1970-01-01T00:00:05 | -6.0 | +| b | h1 | y | 1970-01-01T00:00:10 | -8.0 | ++--------------+------+------+---------------------+-----------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') 2 * {"otel.m"}; + ++--------------+------+------+---------------------+--------------------+ +| service.name | Host | host | ts.time | Float64(2) * v.val | ++--------------+------+------+---------------------+--------------------+ +| a | h1 | x | 1970-01-01T00:00:00 | 2.0 | +| a | h1 | x | 1970-01-01T00:00:05 | 8.0 | +| a | h1 | x | 1970-01-01T00:00:10 | 16.0 | +| a | h2 | x | 1970-01-01T00:00:00 | 4.0 | +| a | h2 | x | 1970-01-01T00:00:05 | 10.0 | +| a | h2 | x | 1970-01-01T00:00:10 | 18.0 | +| b | h1 | y | 1970-01-01T00:00:00 | 6.0 | +| b | h1 | y | 1970-01-01T00:00:05 | 12.0 | +| b | h1 | y | 1970-01-01T00:00:10 | 16.0 | ++--------------+------+------+---------------------+--------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} - 1; + ++--------------+------+------+---------------------+--------------------+ +| service.name | Host | host | ts.time | v.val - Float64(1) | ++--------------+------+------+---------------------+--------------------+ +| a | h1 | x | 1970-01-01T00:00:00 | 0.0 | +| a | h1 | x | 1970-01-01T00:00:05 | 3.0 | +| a | h1 | x | 1970-01-01T00:00:10 | 7.0 | +| a | h2 | x | 1970-01-01T00:00:00 | 1.0 | +| a | h2 | x | 1970-01-01T00:00:05 | 4.0 | +| a | h2 | x | 1970-01-01T00:00:10 | 8.0 | +| b | h1 | y | 1970-01-01T00:00:00 | 2.0 | +| b | h1 | y | 1970-01-01T00:00:05 | 5.0 | +| b | h1 | y | 1970-01-01T00:00:10 | 7.0 | ++--------------+------+------+---------------------+--------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') timestamp({"otel.m"}); + ++---------------------+-------+--------------+------+------+ +| ts.time | value | service.name | Host | host | ++---------------------+-------+--------------+------+------+ +| 1970-01-01T00:00:00 | 0.0 | a | h1 | x | +| 1970-01-01T00:00:00 | 0.0 | a | h2 | x | +| 1970-01-01T00:00:00 | 0.0 | b | h1 | y | +| 1970-01-01T00:00:05 | 5.0 | a | h1 | x | +| 1970-01-01T00:00:05 | 5.0 | a | h2 | x | +| 1970-01-01T00:00:05 | 5.0 | b | h1 | y | +| 1970-01-01T00:00:10 | 10.0 | a | h1 | x | +| 1970-01-01T00:00:10 | 10.0 | a | h2 | x | +| 1970-01-01T00:00:10 | 10.0 | b | h1 | y | ++---------------------+-------+--------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') minute({"otel.m"}); + ++---------------------+-----------------------------------+--------------+------+------+ +| ts.time | date_part(Utf8("minute"),ts.time) | service.name | Host | host | ++---------------------+-----------------------------------+--------------+------+------+ +| 1970-01-01T00:00:00 | 0 | a | h1 | x | +| 1970-01-01T00:00:00 | 0 | a | h2 | x | +| 1970-01-01T00:00:00 | 0 | b | h1 | y | +| 1970-01-01T00:00:05 | 0 | a | h1 | x | +| 1970-01-01T00:00:05 | 0 | a | h2 | x | +| 1970-01-01T00:00:05 | 0 | b | h1 | y | +| 1970-01-01T00:00:10 | 0 | a | h1 | x | +| 1970-01-01T00:00:10 | 0 | a | h2 | x | +| 1970-01-01T00:00:10 | 0 | b | h1 | y | ++---------------------+-----------------------------------+--------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + ++--------------+---------------------+--------------------------------------------------------------+ +| service.name | ts.time | sum(prom_increase(ts.time_range,v.val,ts.time,Int64(10000))) | ++--------------+---------------------+--------------------------------------------------------------+ +| a | 1970-01-01T00:00:05 | 9.0 | +| a | 1970-01-01T00:00:10 | 16.0 | +| b | 1970-01-01T00:00:05 | 6.0 | +| b | 1970-01-01T00:00:10 | 4.0 | ++--------------+---------------------+--------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') scalar(sum({"otel.m"})) * {"otel.m", "Host"="h2"}; + ++--------------+------+------+---------------------+-------------------------------------------+ +| service.name | Host | host | ts.time | lhs.scalar(sum(otel.m.v.val)) * rhs.v.val | ++--------------+------+------+---------------------+-------------------------------------------+ +| a | h2 | x | 1970-01-01T00:00:00 | 12.0 | +| a | h2 | x | 1970-01-01T00:00:05 | 75.0 | +| a | h2 | x | 1970-01-01T00:00:10 | 225.0 | ++--------------+------+------+---------------------+-------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') absent({"otel.m", "service.name"="c"}); + ++---------------------+-------+--------------+ +| ts.time | v.val | service.name | ++---------------------+-------+--------------+ +| 1970-01-01T00:00:00 | 1.0 | c | +| 1970-01-01T00:00:05 | 1.0 | c | +| 1970-01-01T00:00:10 | 1.0 | c | ++---------------------+-------+--------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + ++---------------------+--------------+-------+ +| ts.time | service.name | v.val | ++---------------------+--------------+-------+ +| 1970-01-01T00:00:00 | a | 0.55 | +| 1970-01-01T00:00:05 | a | 0.55 | +| 1970-01-01T00:00:10 | a | 0.775 | ++---------------------+--------------+-------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| LeftSemi Join: otel.m.Host = otel.m.Host, otel.m.host = otel.m.host, otel.m.service.name = otel.m.service.name, otel.m.ts.time = otel.m.ts.time_| +|_|_Aggregate: groupBy=[[otel.m.ts.time, otel.m.service.name, otel.m.Host, otel.m.host, otel.m.v.val]], aggr=[[]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Distinct:_| +|_|_PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.m.ts.time >= TimestampMillisecond(-299999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.ts.time >= TimestampMillisecond(-299999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)]_| +|_| ]]_| +|_|_Projection: otel.m.ts.time, otel.m.service.name, otel.m.Host, otel.m.host_| +|_|_PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Filter: otel.m.Host = Utf8("h1") AND otel.m.ts.time >= TimestampMillisecond(-299999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.Host = Utf8("h1"), otel.m.ts.time >= TimestampMillisecond(-299999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)] | +|_| ]]_| +| physical_plan | HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(Host@2, Host@2), (host@3, host@3), (service.name@1, service.name@1), (ts.time@0, ts.time@0)], NullsEqual: true_| +|_|_CoalescePartitionsExec_| +|_|_ProjectionExec: expr=[ts.time@0 as ts.time, service.name@1 as service.name, Host@2 as Host, host@3 as host]_| +|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivideExec: tags=["service.name", "Host", "host"]_| +|_|_SortExec: expr=[service.name@1 ASC, Host@2 ASC, host@3 ASC, ts.time@0 ASC], preserve_REDACTED +|_|_RepartitionExec: REDACTED +|_|_MergeScanExec: REDACTED +|_|_AggregateExec: mode=FinalPartitioned, gby=[ts.time@0 as ts.time, service.name@1 as service.name, Host@2 as Host, host@3 as host, v.val@4 as v.val], aggr=[]_| +|_|_RepartitionExec: REDACTED +|_|_AggregateExec: mode=Partial, gby=[ts.time@0 as ts.time, service.name@1 as service.name, Host@2 as Host, host@3 as host, v.val@4 as v.val], aggr=[]_| +|_|_RepartitionExec: REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') topk(1, {"otel.m"}); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Projection: otel.m.v.val, otel.m.service.name, otel.m.Host, otel.m.host, otel.m.ts.time_| +|_|_Sort: otel.m.ts.time ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST_| +|_|_Filter: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS Float64) <= Float64(1)_| +|_|_WindowAggr: windowExpr=[[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.m.ts.time >= TimestampMillisecond(-299999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.ts.time >= TimestampMillisecond(-299999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)]_| +|_| ]]_| +| physical_plan | ProjectionExec: expr=[v.val@4 as v.val, service.name@1 as service.name, Host@2 as Host, host@3 as host, ts.time@0 as ts.time]_| +|_|_SortPreservingMergeExec: [ts.time@0 ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 ASC NULLS LAST]_| +|_|_FilterExec: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 AS Float64) <= 1_| +|_|_BoundedWindowAggExec: wdw=[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] | +|_|_SortExec: expr=[ts.time@0 ASC NULLS LAST, v.val@4 DESC, service.name@1 DESC, Host@2 DESC, host@3 DESC], preserve_REDACTED +|_|_RepartitionExec: REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| MergeSort: otel.m.service.name ASC NULLS LAST, otel.m.ts.time ASC NULLS LAST_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Sort: otel.m.service.name ASC NULLS LAST, otel.m.ts.time ASC NULLS LAST_| +|_|_Aggregate: groupBy=[[otel.m.service.name, otel.m.ts.time]], aggr=[[sum(prom_increase(ts.time_range,v.val,ts.time,Int64(10000)))]]_| +|_|_Filter: prom_increase(ts.time_range,v.val,ts.time,Int64(10000)) IS NOT NULL_| +|_|_Projection: otel.m.ts.time, prom_increase(ts.time_range, v.val, otel.m.ts.time, Int64(10000)) AS prom_increase(ts.time_range,v.val,ts.time,Int64(10000)), otel.m.service.name, otel.m.Host, otel.m.host | +|_|_PromRangeManipulate: req range=[5000..10000], interval=[5000], eval range=[10000], time index=[ts.time], values=["v.val"]_| +|_|_PromSeriesNormalize: offset=[0], time index=[ts.time], filter NaN: [true]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.m.ts.time >= TimestampMillisecond(-4999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.ts.time >= TimestampMillisecond(-4999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)]_| +|_| ]]_| +| physical_plan | MergeSortExec: [service.name@0 ASC NULLS LAST, ts.time@1 ASC NULLS LAST]_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| HistogramFold: le=le, field=v.val, quantile=0.5_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "le"]_| +|_|_Sort: otel.h.service.name ASC NULLS FIRST, otel.h.le ASC NULLS FIRST, otel.h.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.h.ts.time >= TimestampMillisecond(-299999, None) AND otel.h.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.h, partial_filters=[otel.h.ts.time >= TimestampMillisecond(-299999, None), otel.h.ts.time <= TimestampMillisecond(10000, None)] | +|_| ]]_| +| physical_plan | HistogramFoldExec: le=@2, field=@3, quantile=0.5_| +|_|_RepartitionExec: REDACTED +|_|_SortExec: expr=[service.name@1 ASC NULLS LAST, ts.time@0 ASC NULLS LAST, TRY_CAST(le@2 AS Float64) ASC NULLS LAST], preserve_REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +DROP TABLE "otel.m"; + +Affected Rows: 0 + +DROP TABLE "otel.h"; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/dotted_column_names.sql b/tests/cases/standalone/common/promql/dotted_column_names.sql new file mode 100644 index 00000000000..e6edc71094b --- /dev/null +++ b/tests/cases/standalone/common/promql/dotted_column_names.sql @@ -0,0 +1,139 @@ +-- https://github.com/GreptimeTeam/greptimedb/issues/9390 +-- Tag, field and time index names containing dots or upper case letters must +-- be resolved as plain column names, not as `relation.column`. `Host` and `host` +-- hold different values, so reading one for the other changes the output. +CREATE TABLE "otel.m" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + "Host" STRING, + host STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", "Host", host) +) PARTITION ON COLUMNS ("service.name") ( + "service.name" < 'b', + "service.name" >= 'b' +); + +INSERT INTO "otel.m" VALUES + (0, 'a', 'h1', 'x', 1), + (0, 'a', 'h2', 'x', 2), + (0, 'b', 'h1', 'y', 3), + (5000, 'a', 'h1', 'x', 4), + (5000, 'a', 'h2', 'x', 5), + (5000, 'b', 'h1', 'y', 6), + (10000, 'a', 'h1', 'x', 8), + (10000, 'a', 'h2', 'x', 9), + (10000, 'b', 'h1', 'y', 8); + +CREATE TABLE "otel.h" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + le STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", le) +); + +INSERT INTO "otel.h" VALUES + (0, 'a', '0.1', 1), + (0, 'a', '1', 3), + (0, 'a', '+Inf', 4), + (10000, 'a', '0.1', 2), + (10000, 'a', '1', 6), + (10000, 'a', '+Inf', 10); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on("service.name") {"otel.m", "service.name"="b"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on(host) {"otel.m", host="y"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless {"otel.m", "Host"="h1"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless on("service.name") {"otel.m", "service.name"="b"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} or {"otel.m", "Host"="h1"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values("v.name", {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values by ("service.name") ("v.name", {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') topk(1, {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') bottomk by ("service.name") (1, {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("service.name") ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("Host") ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by (host) ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum without ("Host") ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') -{"otel.m"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') 2 * {"otel.m"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} - 1; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') timestamp({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') minute({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') scalar(sum({"otel.m"})) * {"otel.m", "Host"="h2"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') absent({"otel.m", "service.name"="c"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') topk(1, {"otel.m"}); + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + +DROP TABLE "otel.m"; + +DROP TABLE "otel.h"; diff --git a/tests/cases/standalone/common/promql/encode_substrait.result b/tests/cases/standalone/common/promql/encode_substrait.result index 230eadb02f7..53bd8a94b1e 100644 --- a/tests/cases/standalone/common/promql/encode_substrait.result +++ b/tests/cases/standalone/common/promql/encode_substrait.result @@ -12,40 +12,42 @@ Affected Rows: 0 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED tql explain (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | PromRangeManipulate: req range=[0..100000], interval=[1000], eval range=[43200000], time index=[ts], values=["prom_increase(ts_range,val,ts,Int64(3600000))"] | -| | PromSeriesDivide: tags=["tag_a", "tag_b"] | -| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | -| | Filter: prom_increase(ts_range,val,ts,Int64(3600000)) IS NOT NULL | -| | Projection: count_total.ts, prom_increase(ts_range, val, count_total.ts, Int64(3600000)) AS prom_increase(ts_range,val,ts,Int64(3600000)), count_total.tag_a, count_total.tag_b | -| | PromRangeManipulate: req range=[-39600000..100000], interval=[3600000], eval range=[3600000], time index=[ts], values=["val"] | -| | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | -| | PromSeriesDivide: tags=["tag_a", "tag_b"] | -| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | -| | Filter: count_total.tag_a = Utf8("ffa") AND count_total.ts >= TimestampMillisecond(-43199999, None) AND count_total.ts <= TimestampMillisecond(100000, None) | -| | TableScan: count_total, partial_filters=[count_total.tag_a = Utf8("ffa"), count_total.ts >= TimestampMillisecond(-43199999, None), count_total.ts <= TimestampMillisecond(100000, None)] | -| | ]] | -| physical_plan | CooperativeExec | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Filter: prom_max_over_time(ts_range,prom_increase(ts_range,val,ts,Int64(3600000))) IS NOT NULL | +| | Projection: count_total.ts, prom_max_over_time(ts_range, prom_increase(ts_range,val,ts,Int64(3600000))) AS prom_max_over_time(ts_range,prom_increase(ts_range,val,ts,Int64(3600000))), count_total.tag_a, count_total.tag_b | +| | PromRangeManipulate: req range=[0..100000], interval=[1000], eval range=[43200000], time index=[ts], values=["prom_increase(ts_range,val,ts,Int64(3600000))"] | +| | PromSeriesDivide: tags=["tag_a", "tag_b"] | +| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | +| | Filter: prom_increase(ts_range,val,ts,Int64(3600000)) IS NOT NULL | +| | Projection: count_total.ts, prom_increase(ts_range, val, count_total.ts, Int64(3600000)) AS prom_increase(ts_range,val,ts,Int64(3600000)), count_total.tag_a, count_total.tag_b | +| | PromRangeManipulate: req range=[-39600000..100000], interval=[3600000], eval range=[3600000], time index=[ts], values=["val"] | +| | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | +| | PromSeriesDivide: tags=["tag_a", "tag_b"] | +| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | +| | Filter: count_total.tag_a = Utf8("ffa") AND count_total.ts >= TimestampMillisecond(-43199999, None) AND count_total.ts <= TimestampMillisecond(100000, None) | +| | TableScan: count_total, partial_filters=[count_total.tag_a = Utf8("ffa"), count_total.ts >= TimestampMillisecond(-43199999, None), count_total.ts <= TimestampMillisecond(100000, None)] | +| | ]] | +| physical_plan | CooperativeExec | | | MergeScanExec: REDACTED -| | | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ tql eval (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); -+----+-----------------------------------------------+-------+-------+----------+ -| ts | prom_increase(ts_range,val,ts,Int64(3600000)) | tag_a | tag_b | ts_range | -+----+-----------------------------------------------+-------+-------+----------+ -+----+-----------------------------------------------+-------+-------+----------+ ++----+----------------------------------------------------------------------------+-------+-------+ +| ts | prom_max_over_time(ts_range,prom_increase(ts_range,val,ts,Int64(3600000))) | tag_a | tag_b | ++----+----------------------------------------------------------------------------+-------+-------+ ++----+----------------------------------------------------------------------------+-------+-------+ drop table count_total; diff --git a/tests/cases/standalone/common/promql/encode_substrait.sql b/tests/cases/standalone/common/promql/encode_substrait.sql index 195c383ebd2..6443f3e1959 100644 --- a/tests/cases/standalone/common/promql/encode_substrait.sql +++ b/tests/cases/standalone/common/promql/encode_substrait.sql @@ -10,13 +10,13 @@ create table count_total ( -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED tql explain (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); tql eval (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); drop table count_total; diff --git a/tests/cases/standalone/common/promql/label.result b/tests/cases/standalone/common/promql/label.result index 112b6261715..f6e169eebe3 100644 --- a/tests/cases/standalone/common/promql/label.result +++ b/tests/cases/standalone/common/promql/label.result @@ -21,9 +21,23 @@ INSERT INTO TABLE test VALUES Affected Rows: 8 -- Missing source labels -- +-- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-"); -Error: 1004(InvalidArguments), Invalid function argument for label_join ++---------------------+-----+----------+-------+------------+ +| ts | val | new_host | host | idc | ++---------------------+-----+----------+-------+------------+ +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:15 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 7 | | host1 | idc4:zone3 | ++---------------------+-----+----------+-------+------------+ -- dst_label is equal to source label -- -- SQLNESS SORT_RESULT 3 1 @@ -63,24 +77,27 @@ TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", "idc", "host" | 1970-01-01T00:00:15 | 7 | idc4:zone3-host1 | idc4:zone3 | +---------------------+-----+------------------+------------+ --- test the empty source label -- --- SQLNESS SORT_RESULT 3 1 +-- Absent labels join as empty strings and keep their separators -- +TQL EVAL (0, 15, '5s') label_join(label_join(vector(1), "a", "", "missing"), "b", "-", "a", "a"); + ++---------------------+----------------+---+---+ +| time | greptime_value | b | a | ++---------------------+----------------+---+---+ +| 1970-01-01T00:00:00 | 1.0 | - | | +| 1970-01-01T00:00:05 | 1.0 | - | | +| 1970-01-01T00:00:10 | 1.0 | - | | +| 1970-01-01T00:00:15 | 1.0 | - | | ++---------------------+----------------+---+---+ + +-- An empty source label name is invalid -- TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", ""); -+---------------------+-----+------+------------+ -| ts | val | host | idc | -+---------------------+-----+------+------------+ -| 1970-01-01T00:00:00 | 1 | | idc1 | -| 1970-01-01T00:00:05 | 1 | | idc1 | -| 1970-01-01T00:00:05 | 3 | | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | | idc1 | -| 1970-01-01T00:00:10 | 3 | | idc2:zone1 | -| 1970-01-01T00:00:10 | 5 | | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | | idc1 | -| 1970-01-01T00:00:15 | 3 | | idc2:zone1 | -| 1970-01-01T00:00:15 | 5 | | idc3:zone2 | -| 1970-01-01T00:00:15 | 7 | | idc4:zone3 | -+---------------------+-----+------+------------+ +Error: 1004(InvalidArguments), Invalid function argument for label_join + +-- Both hosts are joined to the same `host` at 0s, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_join(test, "host", "-", "idc"); + +Error: 3001(EngineExecuteQuery), Execution error: vector cannot contain metrics with the same labelset -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-", "idc", "host"); @@ -114,13 +131,13 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "$2", "idc", +---------------------+-----+---------+-------+------------+ | ts | val | new_idc | host | idc | +---------------------+-----+---------+-------+------------+ -| 1970-01-01T00:00:00 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | | 1970-01-01T00:00:05 | 3 | zone1 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | | 1970-01-01T00:00:10 | 3 | zone1 | host1 | idc2:zone1 | | 1970-01-01T00:00:10 | 5 | zone2 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | | 1970-01-01T00:00:15 | 3 | zone1 | host1 | idc2:zone1 | | 1970-01-01T00:00:15 | 5 | zone2 | host1 | idc3:zone2 | | 1970-01-01T00:00:15 | 7 | zone3 | host1 | idc4:zone3 | @@ -129,20 +146,20 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "$2", "idc", -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "idc99", "idc", "idc2.*"); -+---------------------+-----+------------+-------+------------+ -| ts | val | new_idc | host | idc | -+---------------------+-----+------------+-------+------------+ -| 1970-01-01T00:00:00 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 3 | idc99 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:10 | 3 | idc99 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 5 | idc3:zone2 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:15 | 3 | idc99 | host1 | idc2:zone1 | -| 1970-01-01T00:00:15 | 5 | idc3:zone2 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 7 | idc4:zone3 | host1 | idc4:zone3 | -+---------------------+-----+------------+-------+------------+ ++---------------------+-----+---------+-------+------------+ +| ts | val | new_idc | host | idc | ++---------------------+-----+---------+-------+------------+ +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 3 | idc99 | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 3 | idc99 | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 3 | idc99 | host1 | idc2:zone1 | +| 1970-01-01T00:00:15 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 7 | | host1 | idc4:zone3 | ++---------------------+-----+---------+-------+------------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "new_idc", "$2", "idc", "(.*):(.*)"); @@ -150,26 +167,61 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "new_idc", "$2", "idc", +---------------------+-----+---------+-------+------+ | ts | val | new_idc | host | idc | +---------------------+-----+---------+-------+------+ -| 1970-01-01T00:00:00 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:05 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:05 | 4 | idc2 | host2 | idc2 | -| 1970-01-01T00:00:10 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:10 | 4 | idc2 | host2 | idc2 | -| 1970-01-01T00:00:10 | 6 | idc3 | host2 | idc3 | -| 1970-01-01T00:00:15 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:15 | 4 | idc2 | host2 | idc2 | -| 1970-01-01T00:00:15 | 6 | idc3 | host2 | idc3 | -| 1970-01-01T00:00:15 | 8 | idc4 | host2 | idc4 | +| 1970-01-01T00:00:00 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:05 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:05 | 4 | | host2 | idc2 | +| 1970-01-01T00:00:10 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:10 | 4 | | host2 | idc2 | +| 1970-01-01T00:00:10 | 6 | | host2 | idc3 | +| 1970-01-01T00:00:15 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:15 | 4 | | host2 | idc2 | +| 1970-01-01T00:00:15 | 6 | | host2 | idc3 | +| 1970-01-01T00:00:15 | 8 | | host2 | idc4 | +---------------------+-----+---------+-------+------+ -- dst_label is equal to source label -- -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc", "$2", "idc", "(.*):(.*)"); -Error: 1004(InvalidArguments), vector cannot contain metrics with the same labelset ++---------------------+-----+------+-------+ +| ts | val | idc | host | ++---------------------+-----+------+-------+ +| 1970-01-01T00:00:00 | 2 | idc1 | host2 | +| 1970-01-01T00:00:05 | 2 | idc1 | host2 | +| 1970-01-01T00:00:05 | 4 | idc2 | host2 | +| 1970-01-01T00:00:10 | 2 | idc1 | host2 | +| 1970-01-01T00:00:10 | 4 | idc2 | host2 | +| 1970-01-01T00:00:10 | 6 | idc3 | host2 | +| 1970-01-01T00:00:15 | 2 | idc1 | host2 | +| 1970-01-01T00:00:15 | 4 | idc2 | host2 | +| 1970-01-01T00:00:15 | 6 | idc3 | host2 | +| 1970-01-01T00:00:15 | 8 | idc4 | host2 | ++---------------------+-----+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "idc", "$2", "idc", "(.*):(.*)"); + ++---------------------+-----+-------+-------+ +| ts | val | idc | host | ++---------------------+-----+-------+-------+ +| 1970-01-01T00:00:00 | 1 | idc1 | host1 | +| 1970-01-01T00:00:05 | 1 | idc1 | host1 | +| 1970-01-01T00:00:05 | 3 | zone1 | host1 | +| 1970-01-01T00:00:10 | 1 | idc1 | host1 | +| 1970-01-01T00:00:10 | 3 | zone1 | host1 | +| 1970-01-01T00:00:10 | 5 | zone2 | host1 | +| 1970-01-01T00:00:15 | 1 | idc1 | host1 | +| 1970-01-01T00:00:15 | 3 | zone1 | host1 | +| 1970-01-01T00:00:15 | 5 | zone2 | host1 | +| 1970-01-01T00:00:15 | 7 | zone3 | host1 | ++---------------------+-----+-------+-------+ + +-- Both hosts are rewritten to the same `host`, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_replace(test, "host", "x", "host", ".*"); + +Error: 3001(EngineExecuteQuery), Execution error: vector cannot contain metrics with the same labelset -- test the empty source label -- --- TODO(dennis): we can't remove the label currently -- -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc2", "", "", ""); @@ -221,20 +273,20 @@ TQL EVAL(0, 15, '5s') {__name__="test",host="host1"} + label_replace(vector(1), -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "host2", "host", ""); -+---------------------+-----+-------+------------+ -| ts | val | host | idc | -+---------------------+-----+-------+------------+ -| 1970-01-01T00:00:00 | 1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 3 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | host1 | idc1 | -| 1970-01-01T00:00:10 | 3 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 5 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | host1 | idc1 | -| 1970-01-01T00:00:15 | 3 | host1 | idc2:zone1 | -| 1970-01-01T00:00:15 | 5 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 7 | host1 | idc4:zone3 | -+---------------------+-----+-------+------------+ ++---------------------+-----+-------+-------+------------+ +| ts | val | host2 | host | idc | ++---------------------+-----+-------+-------+------------+ +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:15 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 7 | | host1 | idc4:zone3 | ++---------------------+-----+-------+-------+------------+ -- Empty regex with not existing source label -- SQLNESS SORT_RESULT 3 1 @@ -274,6 +326,23 @@ TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "", "instance", | 1970-01-01T00:00:15 | 7 | host1 | idc4:zone3 | +---------------------+-----+-------+------------+ +-- An empty replacement removes an existing label -- +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{idc="idc1"}, "idc", "", "instance", ""); + ++---------------------+-----+-----+-------+ +| ts | val | idc | host | ++---------------------+-----+-----+-------+ +| 1970-01-01T00:00:00 | 1 | | host1 | +| 1970-01-01T00:00:00 | 2 | | host2 | +| 1970-01-01T00:00:05 | 1 | | host1 | +| 1970-01-01T00:00:05 | 2 | | host2 | +| 1970-01-01T00:00:10 | 1 | | host1 | +| 1970-01-01T00:00:10 | 2 | | host2 | +| 1970-01-01T00:00:15 | 1 | | host1 | +| 1970-01-01T00:00:15 | 2 | | host2 | ++---------------------+-----+-----+-------+ + -- Empty regex and different label value -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') {__name__="test",host="host1"} * label_replace(vector(1), "host", "host2", "host", ""); @@ -307,10 +376,10 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "idc99", "id +---------------------+-----+---------+-------+------+ | ts | val | new_idc | host | idc | +---------------------+-----+---------+-------+------+ -| 1970-01-01T00:00:00 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:10 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:15 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +---------------------+-----+---------+-------+------+ -- SQLNESS SORT_RESULT 3 1 @@ -415,3 +484,52 @@ DROP TABLE test; Affected Rows: 0 +-- Issue 9444 -- +-- Normalizing mixed histogram boundary representations must not be rejected while +-- another label still tells the series apart. +CREATE TABLE demo_bucket ( + ts timestamp(3) time index, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY(instance, le), +); + +Affected Rows: 0 + +INSERT INTO TABLE demo_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +Affected Rows: 7 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (3000, 3000, '5s') label_replace(demo_bucket, "le", "$1", "le", "([0-9]+)[.]0+"); + ++---------------------+------+------+----------+ +| ts | val | le | instance | ++---------------------+------+------+----------+ +| 1970-01-01T00:50:00 | 40.0 | 10 | a | +| 1970-01-01T00:50:00 | 40.0 | 10 | b | +| 1970-01-01T00:50:00 | 50.0 | +Inf | a | +| 1970-01-01T00:50:00 | 50.0 | +Inf | b | +| 1970-01-01T00:50:00 | 50.0 | 100 | a | +| 1970-01-01T00:50:00 | 50.0 | 100 | b | +| 1970-01-01T00:50:00 | 50.0 | 1000 | b | ++---------------------+------+------+----------+ + +-- Aggregating `instance` away first makes the same rewrite collapse two series into +-- one label set, which is a genuine duplicate and must still fail. +TQL EVAL (3000, 3000, '5s') label_replace(sum by (le) (demo_bucket), "le", "$1", "le", "([0-9]+)[.]0+"); + +Error: 3001(EngineExecuteQuery), Execution error: vector cannot contain metrics with the same labelset + +DROP TABLE demo_bucket; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/label.sql b/tests/cases/standalone/common/promql/label.sql index 6a77e9b5214..e269b6a1783 100644 --- a/tests/cases/standalone/common/promql/label.sql +++ b/tests/cases/standalone/common/promql/label.sql @@ -17,6 +17,7 @@ INSERT INTO TABLE test VALUES (15000, 'host2', 'idc4',8); -- Missing source labels -- +-- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-"); -- dst_label is equal to source label -- @@ -27,10 +28,15 @@ TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", "host"); -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", "idc", "host"); --- test the empty source label -- --- SQLNESS SORT_RESULT 3 1 +-- Absent labels join as empty strings and keep their separators -- +TQL EVAL (0, 15, '5s') label_join(label_join(vector(1), "a", "", "missing"), "b", "-", "a", "a"); + +-- An empty source label name is invalid -- TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", ""); +-- Both hosts are joined to the same `host` at 0s, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_join(test, "host", "-", "idc"); + -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-", "idc", "host"); @@ -50,8 +56,13 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "new_idc", "$2", "idc", -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc", "$2", "idc", "(.*):(.*)"); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "idc", "$2", "idc", "(.*):(.*)"); + +-- Both hosts are rewritten to the same `host`, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_replace(test, "host", "x", "host", ".*"); + -- test the empty source label -- --- TODO(dennis): we can't remove the label currently -- -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc2", "", "", ""); @@ -77,6 +88,10 @@ TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "host2", "insta -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "", "instance", ""); +-- An empty replacement removes an existing label -- +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{idc="idc1"}, "idc", "", "instance", ""); + -- Empty regex and different label value -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') {__name__="test",host="host1"} * label_replace(vector(1), "host", "host2", "host", ""); @@ -129,3 +144,32 @@ TQL EVAL (0, 1, '5s') test{job=""}; TQL EVAL (0, 1, '5s') test{job!=""}; DROP TABLE test; + +-- Issue 9444 -- +-- Normalizing mixed histogram boundary representations must not be rejected while +-- another label still tells the series apart. +CREATE TABLE demo_bucket ( + ts timestamp(3) time index, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY(instance, le), +); + +INSERT INTO TABLE demo_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (3000, 3000, '5s') label_replace(demo_bucket, "le", "$1", "le", "([0-9]+)[.]0+"); + +-- Aggregating `instance` away first makes the same rewrite collapse two series into +-- one label set, which is a genuine duplicate and must still fail. +TQL EVAL (3000, 3000, '5s') label_replace(sum by (le) (demo_bucket), "le", "$1", "le", "([0-9]+)[.]0+"); + +DROP TABLE demo_bucket; diff --git a/tests/cases/standalone/common/promql/matching_groups.result b/tests/cases/standalone/common/promql/matching_groups.result new file mode 100644 index 00000000000..04b1588b292 --- /dev/null +++ b/tests/cases/standalone/common/promql/matching_groups.result @@ -0,0 +1,161 @@ +CREATE TABLE matching_groups_left ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); + +Affected Rows: 0 + +CREATE TABLE matching_groups_right ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); + +Affected Rows: 0 + +INSERT INTO matching_groups_left VALUES + ('x', 'a', 0, 12), ('x', 'b', 0, 24), ('y', 'a', 0, 36), + ('', 'a', 0, 48), (NULL, 'a', 0, 60); + +Affected Rows: 5 + +INSERT INTO matching_groups_right VALUES + ('x', 'a', 0, 2), ('x', 'b', 0, 4), ('y', 'a', 0, 100), + ('', 'a', 0, 8), (NULL, 'a', 0, 10); + +Affected Rows: 5 + +-- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); + ++------+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.Float64(8) * greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 24.0 | +| x | b | 1970-01-01T00:00:00 | 48.0 | ++------+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right() (matching_groups_right{host="x"} * 2); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.sum(matching_groups_left.greptime_value) / matching_groups_right.greptime_value * Float64(2) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 9.0 | +| x | b | 1970-01-01T00:00:00 | 4.5 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!="y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | a | 1970-01-01T00:00:00 | 6.0 | +| | a | 1970-01-01T00:00:00 | 6.0 | +| x | a | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | +| y | a | 1970-01-01T00:00:00 | 0.36 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=""} / on(host) group_left max by(host)(matching_groups_right); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | a | 1970-01-01T00:00:00 | 6.0 | +| | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | a | 1970-01-01T00:00:00 | 6.0 | +| | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- bottomk partitions by host when zone is excluded from grouping. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') bottomk without(zone)(1, matching_groups_left) / on(host,zone) matching_groups_right{host="x"}; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- A matcher can leave a global ranking operand without changing its candidates. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') topk(1, matching_groups_left{host="x"}) / on(host,zone) matching_groups_right; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left / on(host,zone) bottomk(1, matching_groups_right{host="x"}); + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- The global winner is host y. Pushing host=x below either topk would change the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk(1, max by(host)(matching_groups_right)); + ++------+------+----+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+----+-------------------------------------------------------------------------------------------------------+ ++------+------+----+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); + ++------+------+----+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+----+-------------------------------------------------------------------------------------------------------+ ++------+------+----+-------------------------------------------------------------------------------------------------------+ + +-- Both one-sides carry two series per `host`, which the matching cannot resolve. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="x"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left max by(host,zone)(matching_groups_right); + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="x"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +DROP TABLE matching_groups_left; + +Affected Rows: 0 + +DROP TABLE matching_groups_right; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/matching_groups.sql b/tests/cases/standalone/common/promql/matching_groups.sql new file mode 100644 index 00000000000..03bdf8eaaa7 --- /dev/null +++ b/tests/cases/standalone/common/promql/matching_groups.sql @@ -0,0 +1,59 @@ +CREATE TABLE matching_groups_left ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); +CREATE TABLE matching_groups_right ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); +INSERT INTO matching_groups_left VALUES + ('x', 'a', 0, 12), ('x', 'b', 0, 24), ('y', 'a', 0, 36), + ('', 'a', 0, 48), (NULL, 'a', 0, 60); +INSERT INTO matching_groups_right VALUES + ('x', 'a', 0, 2), ('x', 'b', 0, 4), ('y', 'a', 0, 100), + ('', 'a', 0, 8), (NULL, 'a', 0, 10); + +-- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right() (matching_groups_right{host="x"} * 2); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!="y"} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=~"x|y"} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=""} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + +-- bottomk partitions by host when zone is excluded from grouping. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') bottomk without(zone)(1, matching_groups_left) / on(host,zone) matching_groups_right{host="x"}; + +-- A matcher can leave a global ranking operand without changing its candidates. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') topk(1, matching_groups_left{host="x"}) / on(host,zone) matching_groups_right; +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left / on(host,zone) bottomk(1, matching_groups_right{host="x"}); + +-- The global winner is host y. Pushing host=x below either topk would change the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk(1, max by(host)(matching_groups_right)); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); + +-- Both one-sides carry two series per `host`, which the matching cannot resolve. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left max by(host,zone)(matching_groups_right); + +DROP TABLE matching_groups_left; +DROP TABLE matching_groups_right; diff --git a/tests/cases/standalone/common/promql/math.result b/tests/cases/standalone/common/promql/math.result index a6271c158cb..e1b884c50f1 100644 --- a/tests/cases/standalone/common/promql/math.result +++ b/tests/cases/standalone/common/promql/math.result @@ -55,15 +55,15 @@ TQL EVAL (0, 20, '5s') pi() + 1; -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 20, '5s') angles{unit="radians"} * pi(); -+---------------------+--------------------+ -| time | angles.val * .pi() | -+---------------------+--------------------+ -| 1970-01-01T00:00:00 | 0.0 | -| 1970-01-01T00:00:05 | 4.934813740258847 | -| 1970-01-01T00:00:10 | 9.869627480517694 | -| 1970-01-01T00:00:15 | 14.80444122077654 | -| 1970-01-01T00:00:20 | 19.73925496103539 | -+---------------------+--------------------+ ++---------+---------------------+--------------------+ +| unit | ts | angles.val * .pi() | ++---------+---------------------+--------------------+ +| radians | 1970-01-01T00:00:00 | 0.0 | +| radians | 1970-01-01T00:00:05 | 4.934813740258847 | +| radians | 1970-01-01T00:00:10 | 9.869627480517694 | +| radians | 1970-01-01T00:00:15 | 14.80444122077654 | +| radians | 1970-01-01T00:00:20 | 19.73925496103539 | ++---------+---------------------+--------------------+ -- Test rad() function (converts degrees to radians) -- SQLNESS SORT_RESULT 3 1 @@ -310,18 +310,18 @@ TQL EVAL (25, 35, '5s') sgn(angles - 42.5); -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 35, '5s') sgn(angles{unit="radians"} - pi()); -+---------------------+----------------------------+ -| time | signum(angles.val - .pi()) | -+---------------------+----------------------------+ -| 1970-01-01T00:00:00 | -1.0 | -| 1970-01-01T00:00:05 | -1.0 | -| 1970-01-01T00:00:10 | 1.0 | -| 1970-01-01T00:00:15 | 1.0 | -| 1970-01-01T00:00:20 | 1.0 | -| 1970-01-01T00:00:25 | 1.0 | -| 1970-01-01T00:00:30 | 1.0 | -| 1970-01-01T00:00:35 | 1.0 | -+---------------------+----------------------------+ ++---------------------+----------------------------+---------+ +| ts | signum(angles.val - .pi()) | unit | ++---------------------+----------------------------+---------+ +| 1970-01-01T00:00:00 | -1.0 | radians | +| 1970-01-01T00:00:05 | -1.0 | radians | +| 1970-01-01T00:00:10 | 1.0 | radians | +| 1970-01-01T00:00:15 | 1.0 | radians | +| 1970-01-01T00:00:20 | 1.0 | radians | +| 1970-01-01T00:00:25 | 1.0 | radians | +| 1970-01-01T00:00:30 | 1.0 | radians | +| 1970-01-01T00:00:35 | 1.0 | radians | ++---------------------+----------------------------+---------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (25, 35, '5s') sgn(angles) * angles; diff --git a/tests/cases/standalone/common/promql/native_time_selection.result b/tests/cases/standalone/common/promql/native_time_selection.result index 2c0e20847cb..945189e68f6 100644 --- a/tests/cases/standalone/common/promql/native_time_selection.result +++ b/tests/cases/standalone/common/promql/native_time_selection.result @@ -613,6 +613,15 @@ TQL EVAL (1, 1, '1s', '1s') native_time_sec offset -500ms; | 11.0 | 1970-01-01T00:00:01 | +------+---------------------+ +-- `timestamp()` reports the timestamp of the selected sample (1s), without the offset applied. +TQL EVAL (1, 1, '1s', '1s') timestamp(native_time_sec offset -500ms); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:00:01 | 1.0 | ++---------------------+-------+ + 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 index ca795b0b3ad..204b047ef23 100644 --- a/tests/cases/standalone/common/promql/native_time_selection.sql +++ b/tests/cases/standalone/common/promql/native_time_selection.sql @@ -219,4 +219,6 @@ 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; +-- `timestamp()` reports the timestamp of the selected sample (1s), without the offset applied. +TQL EVAL (1, 1, '1s', '1s') timestamp(native_time_sec offset -500ms); DROP TABLE native_time_sec; diff --git a/tests/cases/standalone/common/promql/range_vector_result.result b/tests/cases/standalone/common/promql/range_vector_result.result new file mode 100644 index 00000000000..20d60c19480 --- /dev/null +++ b/tests/cases/standalone/common/promql/range_vector_result.result @@ -0,0 +1,143 @@ +-- An instant query of a range-vector expression returns the selected samples +-- with their own timestamps. Expected results are from Prometheus 3. +create table range_vector_result ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +insert into range_vector_result values + (0, 1.0, 'a'), + (10000, 2.0, 'a'), + (20000, 3.0, 'a'), + (30000, 4.0, 'a'), + (0, 10.0, 'b'), + (20000, 30.0, 'b'); + +Affected Rows: 6 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s]; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | +| 1970-01-01T00:00:30 | 4.0 | a | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] offset 10s; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[15s] @ 20; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[30s:10s]; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:10 | 10.0 | b | +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | +| 1970-01-01T00:00:30 | 30.0 | b | +| 1970-01-01T00:00:30 | 4.0 | a | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[35s:10s] offset 5s; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 10.0 | b | +| 1970-01-01T00:00:10 | 10.0 | b | +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | ++---------------------+------+------+ + +tql eval (35, 35, '1s') range_vector_result[4s:10s]; + ++----+-----+------+ +| ts | val | host | ++----+-----+------+ ++----+-----+------+ + +tql eval (0, 30, '10s') range_vector_result[20s]; + +Error: 1004(InvalidArguments), invalid expression type "range vector" for range query, must be Scalar or instant Vector + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] as v; + ++------+------+---------------------+ +| v | host | ts | ++------+------+---------------------+ +| 3.0 | a | 1970-01-01T00:00:20 | +| 30.0 | b | 1970-01-01T00:00:20 | +| 4.0 | a | 1970-01-01T00:00:30 | ++------+------+---------------------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +tql explain (30, 30, '1s') range_vector_result[20s]; + ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | +| | PromSeriesDivide: tags=["host"] | +| | Sort: range_vector_result.host ASC NULLS FIRST, range_vector_result.ts ASC NULLS FIRST | +| | Filter: range_vector_result.ts >= TimestampMillisecond(10001, None) AND range_vector_result.ts <= TimestampMillisecond(30000, None) | +| | TableScan: range_vector_result, partial_filters=[range_vector_result.ts >= TimestampMillisecond(10001, None), range_vector_result.ts <= TimestampMillisecond(30000, None)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +tql eval (30, 30, '1s') range_vector_result[30s:0.0001]; + +Error: 1004(InvalidArguments), Subquery step must be positive + +-- A subquery without a step uses the one-minute default, not the query step. +tql eval (300, 300, '1s') vector(1)[5m:]; + ++---------------------+----------------+ +| time | greptime_value | ++---------------------+----------------+ +| 1970-01-01T00:01:00 | 1.0 | +| 1970-01-01T00:02:00 | 1.0 | +| 1970-01-01T00:03:00 | 1.0 | +| 1970-01-01T00:04:00 | 1.0 | +| 1970-01-01T00:05:00 | 1.0 | ++---------------------+----------------+ + +drop table range_vector_result; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/range_vector_result.sql b/tests/cases/standalone/common/promql/range_vector_result.sql new file mode 100644 index 00000000000..78fdbd81b4f --- /dev/null +++ b/tests/cases/standalone/common/promql/range_vector_result.sql @@ -0,0 +1,49 @@ +-- An instant query of a range-vector expression returns the selected samples +-- with their own timestamps. Expected results are from Prometheus 3. + +create table range_vector_result ( + ts timestamp time index, + val double, + host string primary key +); + +insert into range_vector_result values + (0, 1.0, 'a'), + (10000, 2.0, 'a'), + (20000, 3.0, 'a'), + (30000, 4.0, 'a'), + (0, 10.0, 'b'), + (20000, 30.0, 'b'); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s]; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] offset 10s; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[15s] @ 20; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[30s:10s]; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[35s:10s] offset 5s; + +tql eval (35, 35, '1s') range_vector_result[4s:10s]; + +tql eval (0, 30, '10s') range_vector_result[20s]; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] as v; + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +tql explain (30, 30, '1s') range_vector_result[20s]; + +tql eval (30, 30, '1s') range_vector_result[30s:0.0001]; + +-- A subquery without a step uses the one-minute default, not the query step. +tql eval (300, 300, '1s') vector(1)[5m:]; + +drop table range_vector_result; diff --git a/tests/cases/standalone/common/promql/regex.result b/tests/cases/standalone/common/promql/regex.result index c76b743eafa..c2a10626ff8 100644 --- a/tests/cases/standalone/common/promql/regex.result +++ b/tests/cases/standalone/common/promql/regex.result @@ -146,13 +146,93 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["host"] REDACTED -|_|_|_FilterExec: CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END = REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["CASE WHEN host IS NOT NULL THEN host ELSE Dictionary(UInt32, Utf8(\"\")) END = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host IS NULL OR host = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ +INSERT INTO test VALUES (0, NULL, 2), (0, '', 3), (0, 'ext4', 4), (0, 'tmpfs', 5); + +Affected Rows: 4 + +TQL EVAL (0, 0, '1s') test{host!="tmpfs"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{host!~"tmpfs|vfat"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{host=""}; + ++---------------------+------+-----+ +| ts | host | val | ++---------------------+------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | ++---------------------+------+-----+ + +TQL EVAL (0, 0, '1s') test{host!=""}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | +| 1970-01-01T00:00:00 | tmpfs | 5 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{host=~"ext4|"}; + ++---------------------+------+-----+ +| ts | host | val | ++---------------------+------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | ext4 | 4 | ++---------------------+------+-----+ + +TQL EVAL (0, 0, '1s') test{host!~"ext4|"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | tmpfs | 5 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{missing!="lo"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | +| 1970-01-01T00:00:00 | tmpfs | 5 | ++---------------------+-------------------+-----+ + DROP TABLE test; Affected Rows: 0 diff --git a/tests/cases/standalone/common/promql/regex.sql b/tests/cases/standalone/common/promql/regex.sql index 9804e098e64..80a8f3ed381 100644 --- a/tests/cases/standalone/common/promql/regex.sql +++ b/tests/cases/standalone/common/promql/regex.sql @@ -59,4 +59,14 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".*"}; -- SQLNESS REPLACE (flat_format.*) REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; +INSERT INTO test VALUES (0, NULL, 2), (0, '', 3), (0, 'ext4', 4), (0, 'tmpfs', 5); + +TQL EVAL (0, 0, '1s') test{host!="tmpfs"}; +TQL EVAL (0, 0, '1s') test{host!~"tmpfs|vfat"}; +TQL EVAL (0, 0, '1s') test{host=""}; +TQL EVAL (0, 0, '1s') test{host!=""}; +TQL EVAL (0, 0, '1s') test{host=~"ext4|"}; +TQL EVAL (0, 0, '1s') test{host!~"ext4|"}; +TQL EVAL (0, 0, '1s') test{missing!="lo"}; + DROP TABLE test; diff --git a/tests/cases/standalone/common/promql/set_operation.result b/tests/cases/standalone/common/promql/set_operation.result index 2f0041f1a36..db97fed7e6b 100644 --- a/tests/cases/standalone/common/promql/set_operation.result +++ b/tests/cases/standalone/common/promql/set_operation.result @@ -688,26 +688,26 @@ tql eval (3, 4, '1s') cache_hit_with_null_label / (cache_miss_with_null_label + -- SQLNESS SORT_RESULT 3 1 tql eval (3, 4, '1s') cache_hit_with_null_label / ignoring(null_label) (cache_miss_with_null_label + ignoring(null_label) cache_hit_with_null_label); -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| job | null_label | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| read | | 1970-01-01T00:00:03 | 0.5 | -| read | | 1970-01-01T00:00:04 | 0.75 | -| write | | 1970-01-01T00:00:03 | 0.5 | -| write | | 1970-01-01T00:00:04 | 0.6666666666666666 | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| job | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| read | 1970-01-01T00:00:03 | 0.5 | +| read | 1970-01-01T00:00:04 | 0.75 | +| write | 1970-01-01T00:00:03 | 0.5 | +| write | 1970-01-01T00:00:04 | 0.6666666666666666 | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 tql eval (3, 4, '1s') cache_hit_with_null_label / on(job) (cache_miss_with_null_label + on(job) cache_hit_with_null_label); -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| job | null_label | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| read | | 1970-01-01T00:00:03 | 0.5 | -| read | | 1970-01-01T00:00:04 | 0.75 | -| write | | 1970-01-01T00:00:03 | 0.5 | -| write | | 1970-01-01T00:00:04 | 0.6666666666666666 | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| job | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| read | 1970-01-01T00:00:03 | 0.5 | +| read | 1970-01-01T00:00:04 | 0.75 | +| write | 1970-01-01T00:00:03 | 0.5 | +| write | 1970-01-01T00:00:04 | 0.6666666666666666 | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ drop table cache_hit_with_null_label; diff --git a/tests/cases/standalone/common/promql/simple_histogram.result b/tests/cases/standalone/common/promql/simple_histogram.result index a794a30e8b9..95cc6b6aa30 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.result +++ b/tests/cases/standalone/common/promql/simple_histogram.result @@ -436,3 +436,92 @@ drop table histogram5_bucket; Affected Rows: 0 +-- Issue #9443: string labels stay distinct during aggregation, but equal +-- numeric bucket boundaries must be coalesced before monotonicity repair. +CREATE TABLE equivalent_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY (instance, le) +); + +Affected Rows: 0 + +INSERT INTO equivalent_bounds_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +Affected Rows: 7 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket); + ++--------+---------------------+-----------------------------------+ +| le | ts | sum(equivalent_bounds_bucket.val) | ++--------+---------------------+-----------------------------------+ +| +Inf | 1970-01-01T00:50:00 | 100.0 | +| 10 | 1970-01-01T00:50:00 | 40.0 | +| 10.0 | 1970-01-01T00:50:00 | 40.0 | +| 100 | 1970-01-01T00:50:00 | 50.0 | +| 100.0 | 1970-01-01T00:50:00 | 50.0 | +| 1000.0 | 1970-01-01T00:50:00 | 50.0 | ++--------+---------------------+-----------------------------------+ + +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket)); + ++---------------------+-----------------------------------+ +| ts | sum(equivalent_bounds_bucket.val) | ++---------------------+-----------------------------------+ +| 1970-01-01T00:50:00 | 77.5 | ++---------------------+-----------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket)); + ++----------+---------------------+-----------------------------------+ +| instance | ts | sum(equivalent_bounds_bucket.val) | ++----------+---------------------+-----------------------------------+ +| a | 1970-01-01T00:50:00 | 77.5 | +| b | 1970-01-01T00:50:00 | 77.5 | ++----------+---------------------+-----------------------------------+ + +DROP TABLE equivalent_bounds_bucket; + +Affected Rows: 0 + +-- Equivalent infinity labels belong to the same histogram until the group ends. +CREATE TABLE infinity_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + le STRING, + val DOUBLE, + PRIMARY KEY (le) +); + +Affected Rows: 0 + +INSERT INTO infinity_bounds_bucket VALUES + (3000000, '10', 20), + (3000000, '10.0', 20), + (3000000, '+Inf', 25), + (3000000, '+inf', 25); + +Affected Rows: 4 + +tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket); + ++---------------------+------+ +| ts | val | ++---------------------+------+ +| 1970-01-01T00:50:00 | 6.25 | ++---------------------+------+ + +DROP TABLE infinity_bounds_bucket; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/simple_histogram.sql b/tests/cases/standalone/common/promql/simple_histogram.sql index ec0bf9b1439..6f73325a61f 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.sql +++ b/tests/cases/standalone/common/promql/simple_histogram.sql @@ -253,3 +253,50 @@ insert into histogram5_bucket values tql eval (3000, 3015, '3s') histogram_quantile(0.5, histogram5_bucket); drop table histogram5_bucket; + +-- Issue #9443: string labels stay distinct during aggregation, but equal +-- numeric bucket boundaries must be coalesced before monotonicity repair. +CREATE TABLE equivalent_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY (instance, le) +); + +INSERT INTO equivalent_bounds_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket); + +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket)); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket)); + +DROP TABLE equivalent_bounds_bucket; + +-- Equivalent infinity labels belong to the same histogram until the group ends. +CREATE TABLE infinity_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + le STRING, + val DOUBLE, + PRIMARY KEY (le) +); + +INSERT INTO infinity_bounds_bucket VALUES + (3000000, '10', 20), + (3000000, '10.0', 20), + (3000000, '+Inf', 25), + (3000000, '+inf', 25); + +tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket); + +DROP TABLE infinity_bounds_bucket; diff --git a/tests/cases/standalone/common/promql/sub_millisecond_interval.result b/tests/cases/standalone/common/promql/sub_millisecond_interval.result new file mode 100644 index 00000000000..6461c8e0b0d --- /dev/null +++ b/tests/cases/standalone/common/promql/sub_millisecond_interval.result @@ -0,0 +1,77 @@ +-- The last two 'sub_ms' samples are 400ns apart and share one millisecond, which +-- irate and idelta cannot resolve, so they return no point instead of inf. +-- The 'one_ms' samples are 1ms apart and still produce a rate. +CREATE TABLE sub_ms_interval ( + job STRING, + `value` DOUBLE, + ts TIMESTAMP(9) TIME INDEX, + PRIMARY KEY(job) +); + +Affected Rows: 0 + +INSERT INTO sub_ms_interval VALUES + ('sub_ms', 999, 999999900), + ('sub_ms', 10, 999999100), + ('sub_ms', 13, 999999500), + ('one_ms', 10, 998000000), + ('one_ms', 13, 999000000); + +Affected Rows: 5 + +SELECT * FROM sub_ms_interval ORDER BY job, ts; + ++--------+-------+-------------------------------+ +| job | value | ts | ++--------+-------+-------------------------------+ +| one_ms | 10.0 | 1970-01-01T00:00:00.998 | +| one_ms | 13.0 | 1970-01-01T00:00:00.999 | +| sub_ms | 10.0 | 1970-01-01T00:00:00.999999100 | +| sub_ms | 13.0 | 1970-01-01T00:00:00.999999500 | +| sub_ms | 999.0 | 1970-01-01T00:00:00.999999900 | ++--------+-------+-------------------------------+ + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + ++---------------------+----------------------------+--------+ +| ts | prom_irate(ts_range,value) | job | ++---------------------+----------------------------+--------+ +| 1970-01-01T00:00:01 | 3000.0 | one_ms | ++---------------------+----------------------------+--------+ + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + ++---------------------+-----------------------------+--------+ +| ts | prom_idelta(ts_range,value) | job | ++---------------------+-----------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | one_ms | ++---------------------+-----------------------------+--------+ + +ADMIN FLUSH_TABLE('sub_ms_interval'); + ++--------------------------------------+ +| ADMIN FLUSH_TABLE('sub_ms_interval') | ++--------------------------------------+ +| 0 | ++--------------------------------------+ + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + ++---------------------+----------------------------+--------+ +| ts | prom_irate(ts_range,value) | job | ++---------------------+----------------------------+--------+ +| 1970-01-01T00:00:01 | 3000.0 | one_ms | ++---------------------+----------------------------+--------+ + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + ++---------------------+-----------------------------+--------+ +| ts | prom_idelta(ts_range,value) | job | ++---------------------+-----------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | one_ms | ++---------------------+-----------------------------+--------+ + +DROP TABLE sub_ms_interval; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/sub_millisecond_interval.sql b/tests/cases/standalone/common/promql/sub_millisecond_interval.sql new file mode 100644 index 00000000000..d643479390d --- /dev/null +++ b/tests/cases/standalone/common/promql/sub_millisecond_interval.sql @@ -0,0 +1,30 @@ +-- The last two 'sub_ms' samples are 400ns apart and share one millisecond, which +-- irate and idelta cannot resolve, so they return no point instead of inf. +-- The 'one_ms' samples are 1ms apart and still produce a rate. +CREATE TABLE sub_ms_interval ( + job STRING, + `value` DOUBLE, + ts TIMESTAMP(9) TIME INDEX, + PRIMARY KEY(job) +); + +INSERT INTO sub_ms_interval VALUES + ('sub_ms', 999, 999999900), + ('sub_ms', 10, 999999100), + ('sub_ms', 13, 999999500), + ('one_ms', 10, 998000000), + ('one_ms', 13, 999000000); + +SELECT * FROM sub_ms_interval ORDER BY job, ts; + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + +ADMIN FLUSH_TABLE('sub_ms_interval'); + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + +DROP TABLE sub_ms_interval; diff --git a/tests/cases/standalone/common/promql/subquery.result b/tests/cases/standalone/common/promql/subquery.result index 12e65c4310c..1be8073a541 100644 --- a/tests/cases/standalone/common/promql/subquery.result +++ b/tests/cases/standalone/common/promql/subquery.result @@ -63,3 +63,270 @@ drop table metric_total; Affected Rows: 0 +-- Offset on a subquery shifts the subquery's own evaluation window back by the offset. +-- Reference: Prometheus `evaluator.subqueryTimeRange` (promql/engine.go). +-- The offset cases stay on the subquery step grid so they match Prometheus directly. +create table subquery_offset_total ( + ts timestamp time index, + host string primary key, + val double, +); + +Affected Rows: 0 + +insert into subquery_offset_total values + (0, 'a', 1), + (10000, 'a', 2), + (20000, 'a', 3), + (30000, 'a', 4), + (40000, 'a', 5), + (50000, 'a', 6), + (60000, 'a', 7); + +Affected Rows: 7 + +-- baseline: no offset at t=60 covers the 10s subquery points in (40s, 60s] -> 6 + 7 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 13.0 | a | ++---------------------+----------------------------------+------+ + +-- the same subquery evaluated at t=30 -> 3 + 4 +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- `offset 30s` at t=60 must equal the un-offset subquery at t=30 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- a negative offset looks ahead of the evaluation time +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s] offset -30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 13.0 | a | ++---------------------+----------------------------------+------+ + +-- an offset on the inner selector composes additively with the subquery offset: Prometheus +-- `subqueryTimes` accumulates "the sum of offsets and ranges of all subqueries in the path", +-- and the inner selector subtracts its own offset from the already shifted step timestamps. +-- 20s + 10s therefore behaves like the un-offset subquery at t=30. +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 10s)[20s:10s] offset 20s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- ... and the inner offset alone accounts for the same total shift +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 30s)[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- `predict_linear` predicts from the evaluation time: 4 + 0.1 * 30 = 7 and 7 - 0.1 * 30 = 4 +tql eval (60, 60, '1s') predict_linear(subquery_offset_total[20s:10s] offset 30s, 0); + ++---------------------+----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(0)) | host | ++---------------------+----------------------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------------------+------+ + +tql eval (30, 30, '1s') predict_linear(subquery_offset_total[20s:10s] offset -30s, 0); + ++---------------------+----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(0)) | host | ++---------------------+----------------------------------------------+------+ +| 1970-01-01T00:00:30 | 4.0 | a | ++---------------------+----------------------------------------------+------+ + +-- Offset results match the shifted evaluation while retaining outer timestamps. +-- Verified against Prometheus v3.14.0. +tql eval (30, 60, '30s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 1.0 | a | +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +tql eval (0, 30, '30s') sum_over_time(subquery_offset_total[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:30 | 7.0 | a | ++---------------------+----------------------------------+------+ + +drop table subquery_offset_total; + +Affected Rows: 0 + +-- Off-grid samples expose absolute step anchoring. +create table subquery_offset_offgrid ( + ts timestamp time index, + host string primary key, + val double, +); + +Affected Rows: 0 + +insert into subquery_offset_offgrid values + (5000, 'a', 1), + (15000, 'a', 2), + (25000, 'a', 3), + (35000, 'a', 4), + (45000, 'a', 5), + (55000, 'a', 6); + +Affected Rows: 6 + +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 3.0 | a | ++---------------------+----------------------------------+------+ + +-- Offset moves the window without re-phasing its step points. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 5.0 | a | ++---------------------+----------------------------------+------+ + +-- The window starts before the first sample. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 15s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 1.0 | a | ++---------------------+----------------------------------+------+ + +tql eval (15, 15, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset -5s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:15 | 3.0 | a | ++---------------------+----------------------------------+------+ + +-- Multi-step evaluation uses the non-aligned parent timeline. +tql eval (35, 65, '30s') sum_over_time(subquery_offset_offgrid[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:35 | 5.0 | a | +| 1970-01-01T00:01:05 | 11.0 | a | ++---------------------+----------------------------------+------+ + +-- An evaluation instant whose window reaches no sample has no row (30s), unlike 60s -> 2 + 3 +tql eval (30, 60, '30s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 5.0 | a | ++---------------------+----------------------------------+------+ + +-- No child point yields no row; the enclosing fallback still fires. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + ++----+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++----+----------------------------------+------+ ++----+----------------------------------+------+ + +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1); + ++---------------------+------+----------------------------------+ +| ts | host | prom_sum_over_time(ts_range,val) | ++---------------------+------+----------------------------------+ +| 1970-01-01T00:00:35 | | 1.0 | ++---------------------+------+----------------------------------+ + +-- A negative offset makes the short window reach a child point. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s] offset -5s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:35 | 4.0 | a | ++---------------------+----------------------------------+------+ + +-- An empty outer window suppresses inner results, but not an outer fallback. +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]); + ++----+---------------------------------------------------------------+------+ +| ts | prom_sum_over_time(ts_range,prom_sum_over_time(ts_range,val)) | host | ++----+---------------------------------------------------------------+------+ ++----+---------------------------------------------------------------+------+ + +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]) or vector(1); + ++---------------------+------+---------------------------------------------------------------+ +| ts | host | prom_sum_over_time(ts_range,prom_sum_over_time(ts_range,val)) | ++---------------------+------+---------------------------------------------------------------+ +| 1970-01-01T00:00:35 | | 1.0 | ++---------------------+------+---------------------------------------------------------------+ + +-- An aligned instant has a child point in the same short range. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[5s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 3.0 | a | ++---------------------+----------------------------------+------+ + +-- A longer range covers child steps on the off-grid parent timeline. +tql eval (35, 65, '10s') sum_over_time(subquery_offset_offgrid[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:35 | 5.0 | a | +| 1970-01-01T00:00:45 | 7.0 | a | +| 1970-01-01T00:00:55 | 9.0 | a | +| 1970-01-01T00:01:05 | 11.0 | a | ++---------------------+----------------------------------+------+ + +-- Zero step parses here but is invalid for planning. +tql eval (60, 60, '1s') sum_over_time(subquery_offset_offgrid[5s:0]); + +Error: 1004(InvalidArguments), Zero range in range selector + +tql eval (44, 35, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +Error: 1004(InvalidArguments), The end time must be greater than start time, start: 44000, end: 35000 + +drop table subquery_offset_offgrid; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/subquery.sql b/tests/cases/standalone/common/promql/subquery.sql index 95215b8488b..c44e5163b8d 100644 --- a/tests/cases/standalone/common/promql/subquery.sql +++ b/tests/cases/standalone/common/promql/subquery.sql @@ -20,3 +20,112 @@ tql eval (10, 10, '1s') rate(metric_total[20s:10s]); tql eval (20, 20, '1s') rate(metric_total[20s:5s]); drop table metric_total; + +-- Offset on a subquery shifts the subquery's own evaluation window back by the offset. +-- Reference: Prometheus `evaluator.subqueryTimeRange` (promql/engine.go). +-- The offset cases stay on the subquery step grid so they match Prometheus directly. +create table subquery_offset_total ( + ts timestamp time index, + host string primary key, + val double, +); + +insert into subquery_offset_total values + (0, 'a', 1), + (10000, 'a', 2), + (20000, 'a', 3), + (30000, 'a', 4), + (40000, 'a', 5), + (50000, 'a', 6), + (60000, 'a', 7); + +-- baseline: no offset at t=60 covers the 10s subquery points in (40s, 60s] -> 6 + 7 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s]); + +-- the same subquery evaluated at t=30 -> 3 + 4 +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s]); + +-- `offset 30s` at t=60 must equal the un-offset subquery at t=30 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + +-- a negative offset looks ahead of the evaluation time +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s] offset -30s); + +-- an offset on the inner selector composes additively with the subquery offset: Prometheus +-- `subqueryTimes` accumulates "the sum of offsets and ranges of all subqueries in the path", +-- and the inner selector subtracts its own offset from the already shifted step timestamps. +-- 20s + 10s therefore behaves like the un-offset subquery at t=30. +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 10s)[20s:10s] offset 20s); + +-- ... and the inner offset alone accounts for the same total shift +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 30s)[20s:10s]); + +-- `predict_linear` predicts from the evaluation time: 4 + 0.1 * 30 = 7 and 7 - 0.1 * 30 = 4 +tql eval (60, 60, '1s') predict_linear(subquery_offset_total[20s:10s] offset 30s, 0); + +tql eval (30, 30, '1s') predict_linear(subquery_offset_total[20s:10s] offset -30s, 0); + +-- Offset results match the shifted evaluation while retaining outer timestamps. +-- Verified against Prometheus v3.14.0. +tql eval (30, 60, '30s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + +tql eval (0, 30, '30s') sum_over_time(subquery_offset_total[20s:10s]); + +drop table subquery_offset_total; + +-- Off-grid samples expose absolute step anchoring. +create table subquery_offset_offgrid ( + ts timestamp time index, + host string primary key, + val double, +); + +insert into subquery_offset_offgrid values + (5000, 'a', 1), + (15000, 'a', 2), + (25000, 'a', 3), + (35000, 'a', 4), + (45000, 'a', 5), + (55000, 'a', 6); + +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s); + +-- Offset moves the window without re-phasing its step points. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s]); + +-- The window starts before the first sample. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 15s); + +tql eval (15, 15, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset -5s); + +-- Multi-step evaluation uses the non-aligned parent timeline. +tql eval (35, 65, '30s') sum_over_time(subquery_offset_offgrid[20s:10s]); + +-- An evaluation instant whose window reaches no sample has no row (30s), unlike 60s -> 2 + 3 +tql eval (30, 60, '30s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 30s); + +-- No child point yields no row; the enclosing fallback still fires. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1); + +-- A negative offset makes the short window reach a child point. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s] offset -5s); + +-- An empty outer window suppresses inner results, but not an outer fallback. +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]); + +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]) or vector(1); + +-- An aligned instant has a child point in the same short range. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +-- A longer range covers child steps on the off-grid parent timeline. +tql eval (35, 65, '10s') sum_over_time(subquery_offset_offgrid[20s:10s]); + +-- Zero step parses here but is invalid for planning. +tql eval (60, 60, '1s') sum_over_time(subquery_offset_offgrid[5s:0]); + +tql eval (44, 35, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +drop table subquery_offset_offgrid; diff --git a/tests/cases/standalone/common/promql/time_fn.result b/tests/cases/standalone/common/promql/time_fn.result index 25e0cd11030..f947f2fb726 100644 --- a/tests/cases/standalone/common/promql/time_fn.result +++ b/tests/cases/standalone/common/promql/time_fn.result @@ -150,7 +150,7 @@ tql eval (1, 2, '1s') time() == bool metrics; tql eval (1, 2, '1s') metrics + time(); +---------------------+-------------------------------------+ -| time | metrics.val + .time / Float64(1000) | +| ts | metrics.val + .time / Float64(1000) | +---------------------+-------------------------------------+ | 1970-01-01T00:00:01 | 2.0 | | 1970-01-01T00:00:02 | 4.0 | @@ -170,7 +170,7 @@ tql eval (1, 2, '1s') metrics == time(); tql eval (1, 2, '1s') metrics == bool time(); +---------------------+-------------------------------------+ -| time | metrics.val = .time / Float64(1000) | +| ts | metrics.val = .time / Float64(1000) | +---------------------+-------------------------------------+ | 1970-01-01T00:00:01 | 1.0 | | 1970-01-01T00:00:02 | 1.0 | diff --git a/tests/cases/standalone/common/promql/timestamp_fn.result b/tests/cases/standalone/common/promql/timestamp_fn.result index 1f5682649a4..32f5ba11cfb 100644 --- a/tests/cases/standalone/common/promql/timestamp_fn.result +++ b/tests/cases/standalone/common/promql/timestamp_fn.result @@ -95,7 +95,7 @@ tql eval (0, 60, '30s') timestamp(timestamp_test) > bool 30; tql eval (0, 60, '30s') timestamp(timestamp_test) - time(); +---------------------+----------------------------------------------+ -| time | timestamp_test.value - .time / Float64(1000) | +| ts | timestamp_test.value - .time / Float64(1000) | +---------------------+----------------------------------------------+ | 1970-01-01T00:00:00 | 0.0 | | 1970-01-01T00:00:30 | -29.0 | @@ -113,17 +113,17 @@ tql eval (0, 60, '30s') abs(timestamp(timestamp_test) - avg(timestamp(timestamp_ -- Test Issue 6707 tql eval timestamp(demo_memory_usage_bytes * 1); -+------+--------------------+ -| time | value * Float64(1) | -+------+--------------------+ -+------+--------------------+ ++------+-------+ +| time | value | ++------+-------+ ++------+-------+ tql eval timestamp(-demo_memory_usage_bytes); -+------+-----------+ -| time | (- value) | -+------+-----------+ -+------+-----------+ ++------+-------+ +| time | value | ++------+-------+ ++------+-------+ tql eval (0, 60, '30s') timestamp(timestamp_test) == 60; @@ -133,6 +133,79 @@ tql eval (0, 60, '30s') timestamp(timestamp_test) == 60; | 1970-01-01T00:01:00 | 60.0 | +---------------------+-------+ +-- timestamp() reports the timestamp of the selected sample, not the evaluation time +tql eval (60, 60, '1s') timestamp(timestamp_test offset 30s); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:01:00 | 1.0 | ++---------------------+-------+ + +tql eval (60, 60, '1s') timestamp((timestamp_test offset 30s)); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:01:00 | 1.0 | ++---------------------+-------+ + +-- The result of a function is sampled at the evaluation time +tql eval (90, 90, '1s') timestamp(abs(timestamp_test)); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:01:30 | 90.0 | ++---------------------+-------+ + +-- A row whose fields are all NULL holds no sample and gets no timestamp; a row with any +-- non-NULL field does. +create table timestamp_null (ts timestamp time index, host string primary key, a double, b double); + +Affected Rows: 0 + +insert into timestamp_null values + (0, 'none', NULL, NULL), + (0, 'only_a', 1.0, NULL), + (0, 'only_b', NULL, 2.0); + +Affected Rows: 3 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(timestamp_null); + ++---------------------+-------+--------+ +| ts | value | host | ++---------------------+-------+--------+ +| 1970-01-01T00:01:00 | 0.0 | only_a | +| 1970-01-01T00:01:00 | 0.0 | only_b | ++---------------------+-------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(-timestamp_null); + ++---------------------+-------+--------+ +| ts | value | host | ++---------------------+-------+--------+ +| 1970-01-01T00:01:00 | 60.0 | only_a | +| 1970-01-01T00:01:00 | 60.0 | only_b | ++---------------------+-------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(abs(timestamp_null)); + ++---------------------+-------+--------+ +| ts | value | host | ++---------------------+-------+--------+ +| 1970-01-01T00:01:00 | 60.0 | only_a | +| 1970-01-01T00:01:00 | 60.0 | only_b | ++---------------------+-------+--------+ + +drop table timestamp_null; + +Affected Rows: 0 + -- Test timestamp() with multiple metrics create table timestamp_test2 (ts timestamp time index, val double); diff --git a/tests/cases/standalone/common/promql/timestamp_fn.sql b/tests/cases/standalone/common/promql/timestamp_fn.sql index 123bdeb5881..cc307007e4b 100644 --- a/tests/cases/standalone/common/promql/timestamp_fn.sql +++ b/tests/cases/standalone/common/promql/timestamp_fn.sql @@ -43,6 +43,34 @@ tql eval timestamp(-demo_memory_usage_bytes); tql eval (0, 60, '30s') timestamp(timestamp_test) == 60; +-- timestamp() reports the timestamp of the selected sample, not the evaluation time +tql eval (60, 60, '1s') timestamp(timestamp_test offset 30s); + +tql eval (60, 60, '1s') timestamp((timestamp_test offset 30s)); + +-- The result of a function is sampled at the evaluation time +tql eval (90, 90, '1s') timestamp(abs(timestamp_test)); + +-- A row whose fields are all NULL holds no sample and gets no timestamp; a row with any +-- non-NULL field does. +create table timestamp_null (ts timestamp time index, host string primary key, a double, b double); + +insert into timestamp_null values + (0, 'none', NULL, NULL), + (0, 'only_a', 1.0, NULL), + (0, 'only_b', NULL, 2.0); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(timestamp_null); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(-timestamp_null); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(abs(timestamp_null)); + +drop table timestamp_null; + -- Test timestamp() with multiple metrics create table timestamp_test2 (ts timestamp time index, val double); diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index d80d29ce000..fc4bc6f01c2 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -127,13 +127,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -167,13 +167,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -210,19 +210,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (ts |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -243,25 +243,32 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[host@0 as host, job@1 as job, ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 / greptime_value@5 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@2), (ts@2, ts@4)], projection=[host@4, job@5, ts@7, __tsid@6, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED +| 0_| 0_|_ProjectionExec: expr=[job@0 as job, ts@1 as ts, greptime_value@2 / greptime_value@3 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@4, job@1), projection=[job@1, ts@3, greptime_value@0, greptime_value@2] REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@1 as job, greptime_value@3 as greptime_value, ts@4 as ts, __promql_match_group_count@5 as __promql_match_group_count] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@2)], projection=[greptime_value@0, job@1, ts@2, greptime_value@3, ts@5], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_RepartitionExec: partitioning=Hash([job@2, ts@4],REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@3, job@1), projection=[greptime_value@0, job@1, ts@2] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -282,25 +289,32 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[job@0 as job, ts@1 as ts, __tsid@2 as __tsid, greptime_value@3 / greptime_value@4 as tsid_binary_join_left.greptime_value / tsid_binary_join_right_by_job.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@3)], projection=[job@4, ts@6, __tsid@5, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED +| 0_| 0_|_ProjectionExec: expr=[job@0 as job, ts@1 as ts, greptime_value@2 / greptime_value@3 as tsid_binary_join_left.greptime_value / tsid_binary_join_right_by_job.greptime_value] REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@4, job@1), projection=[job@1, ts@3, greptime_value@0, greptime_value@2] REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@1 as job, greptime_value@3 as greptime_value, ts@4 as ts, __promql_match_group_count@5 as __promql_match_group_count] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED +|_|_|_CoalescePartitionsExec REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@2)], projection=[greptime_value@0, job@1, ts@2, greptime_value@3, ts@5], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@3],REDACTED +|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@1 as job, ts@3 as ts] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -333,13 +347,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -372,13 +386,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -416,19 +430,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -471,25 +485,25 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -527,7 +541,7 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 as greptime_value, host@5 as host, job@6 as job, CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE_END as __promql_or_match_0, CASE WHEN __common_expr_2@1 IS NOT NULL THEN __common_expr_2@1 ELSE_END as __promql_or_match_1] REDACTED |_|_|_ProjectionExec: expr=[CAST(host@1 AS Utf8) as __common_expr_1, CAST(job@2 AS Utf8) as __common_expr_2, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED @@ -535,13 +549,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -566,13 +580,16 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@0, ts@4)], projection=[host@4, job@5, ts@7, __tsid@6, tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value@2, greptime_value@3], NullsEqual: true REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, greptime_value@2 / greptime_value@3 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@3)], projection=[ts@6, __tsid@5, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED -|_|_|_RepartitionExec: partitioning=Hash([REDACTED -|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED -|_|_|_MergeScanExec: REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@3, ts@2)], projection=[ts@6, __tsid@2, greptime_value@0, greptime_value@4], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@3, job@1), projection=[greptime_value@0, job@1, ts@2] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=Hash([REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| @@ -580,23 +597,37 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ +-- A group modifier that keeps the many side's whole tag set keeps its `__tsid` too, so the +-- enclosing operation matches on it. Same rows as the label-based plan. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; + ++-------+------+---------------------+--------------------------------------------------------------------------------------------------------------------------------------------+ +| host | job | ts | tsid_binary_join_right.tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value / tsid_binary_join_left.greptime_value | ++-------+------+---------------------+--------------------------------------------------------------------------------------------------------------------------------------------+ +| host1 | job1 | 1970-01-01T00:00:00 | 0.3333333333333333 | +| host1 | job1 | 1970-01-01T00:00:05 | 0.2 | +| host2 | job2 | 1970-01-01T00:00:00 | 0.16666666666666666 | +| host2 | job2 | 1970-01-01T00:00:05 | 0.14285714285714285 | ++-------+------+---------------------+--------------------------------------------------------------------------------------------------------------------------------------------+ + -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql b/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql index 7a3bc897c75..627680330aa 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql @@ -226,6 +226,11 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; +-- A group modifier that keeps the many side's whole tag set keeps its `__tsid` too, so the +-- enclosing operation matches on it. Same rows as the label-based plan. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; + -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index e50fa43f10e..10b31cac1e1 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -118,7 +118,7 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, le@4 as le, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/vector_matching.result b/tests/cases/standalone/common/promql/vector_matching.result new file mode 100644 index 00000000000..5382e94967c --- /dev/null +++ b/tests/cases/standalone/common/promql/vector_matching.result @@ -0,0 +1,406 @@ +CREATE TABLE counter_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +Affected Rows: 0 + +CREATE TABLE gauge_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +Affected Rows: 0 + +INSERT INTO counter_metric VALUES + ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20), ('host2', 'eth0', 0, 30); + +Affected Rows: 3 + +INSERT INTO gauge_metric VALUES + ('host1', 'eth0', 0, 2), ('host1', 'eth1', 0, 4), ('host2', 'eth0', 0, 5); + +Affected Rows: 3 + +-- Matching on the whole tag set keeps it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / gauge_metric; + ++-------+--------+---------------------+-------------------------------------------------------------+ +| host | device | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+--------+---------------------+-------------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 6.0 | ++-------+--------+---------------------+-------------------------------------------------------------+ + +-- `on(...)` reduces the result to the matching labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) gauge_metric{device="eth0"}; + ++-------+---------------------+-------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+---------------------+-------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 5.0 | +| host2 | 1970-01-01T00:00:00 | 6.0 | ++-------+---------------------+-------------------------------------------------------------+ + +-- `ignoring(...)` drops the ignored labels, here from operands that don't even share them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / ignoring(device) gauge_metric{device="eth1"}; + ++-------+---------------------+-------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+---------------------+-------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 2.5 | ++-------+---------------------+-------------------------------------------------------------+ + +-- `on()` matches everything against everything, so both sides must hold a single series. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{host="host1",device="eth0"} / on() gauge_metric{host="host1",device="eth0"}; + ++---------------------+-------------------------------------------------------------+ +| ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++---------------------+-------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 5.0 | ++---------------------+-------------------------------------------------------------+ + +-- A filtering comparison keeps the left sample values and the derived labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > on(host) gauge_metric{device="eth0"}; + ++---------------------+----------------+-------+ +| ts | greptime_value | host | ++---------------------+----------------+-------+ +| 1970-01-01T00:00:00 | 10.0 | host1 | +| 1970-01-01T00:00:00 | 30.0 | host2 | ++---------------------+----------------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > bool on(host) gauge_metric{device="eth0"}; + ++-------+---------------------+-------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value > gauge_metric.greptime_value | ++-------+---------------------+-------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 1.0 | +| host2 | 1970-01-01T00:00:00 | 1.0 | ++-------+---------------------+-------------------------------------------------------------+ + +-- `host1` carries two devices on both sides: the match group repeats on the one side. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric; + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +-- The one side is unique here, the many side is not, and no group modifier makes it explicit. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; + +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) + +-- `group_left` takes the result labels from the many side. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left gauge_metric{device="eth0"}; + ++-------+--------+---------------------+-------------------------------------------------------------+ +| host | device | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+--------+---------------------+-------------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 6.0 | ++-------+--------+---------------------+-------------------------------------------------------------+ + +-- `group_right` swaps the sides. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_right gauge_metric; + ++-------+--------+---------------------+-------------------------------------------------------------+ +| host | device | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+--------+---------------------+-------------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 2.5 | +| host2 | eth0 | 1970-01-01T00:00:00 | 6.0 | ++-------+--------+---------------------+-------------------------------------------------------------+ + +-- `group_right` makes the left operand the one side, which `host1` breaks. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_right gauge_metric{device="eth0"}; + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the left hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +-- `group_left(device)` copies `device` from the one side, so both `host1` rows end up with the +-- same labels. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left(device) gauge_metric{device="eth0"}; + +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1", device="eth0"}: grouping labels must ensure unique matches + +-- An included label the one side doesn't carry is dropped from the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_left(device) sum by(host)(gauge_metric); + ++-------+---------------------+-------------------------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value / gauge_metric.sum(gauge_metric.greptime_value) | ++-------+---------------------+-------------------------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 1.6666666666666667 | +| host2 | 1970-01-01T00:00:00 | 6.0 | ++-------+---------------------+-------------------------------------------------------------------------------+ + +DROP TABLE counter_metric; + +Affected Rows: 0 + +DROP TABLE gauge_metric; + +Affected Rows: 0 + +-- A label can carry the name the cardinality check generates for its row count. +CREATE TABLE collide_left ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +Affected Rows: 0 + +CREATE TABLE collide_right ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +Affected Rows: 0 + +INSERT INTO collide_left VALUES ('host1', 'a', 0, 10); + +Affected Rows: 1 + +INSERT INTO collide_right VALUES ('host1', 'b', 0, 2); + +Affected Rows: 1 + +TQL EVAL (0, 0, '5s') collide_left / on(host) collide_right; + ++-------+---------------------+------------------------------------------------------------+ +| host | ts | collide_left.greptime_value / collide_right.greptime_value | ++-------+---------------------+------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 5.0 | ++-------+---------------------+------------------------------------------------------------+ + +DROP TABLE collide_left; + +Affected Rows: 0 + +DROP TABLE collide_right; + +Affected Rows: 0 + +-- Partitioning on a column outside the match keys puts one match group in several regions. The +-- cardinality check counts the merged data, so it must still see the duplicate. +CREATE TABLE spread_left ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +) +PARTITION ON COLUMNS (device) ( + device < 'eth1', + device >= 'eth1' +); + +Affected Rows: 0 + +CREATE TABLE spread_right ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +INSERT INTO spread_left VALUES ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20); + +Affected Rows: 2 + +INSERT INTO spread_right VALUES ('host1', 0, 2); + +Affected Rows: 1 + +TQL EVAL (0, 0, '5s') spread_left / on(host) spread_right; + +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) + +TQL EVAL (0, 0, '5s') spread_right / on(host) group_left spread_left; + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +DROP TABLE spread_left; + +Affected Rows: 0 + +DROP TABLE spread_right; + +Affected Rows: 0 + +-- An outer aggregate groups on the labels the operand has left, not on the ones the scan +-- underneath could still produce. +CREATE TABLE outer_agg_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +Affected Rows: 0 + +CREATE TABLE outer_agg_b ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +INSERT INTO outer_agg_a VALUES ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +Affected Rows: 2 + +INSERT INTO outer_agg_b VALUES ('h1', 0, 2), ('h2', 0, 4); + +Affected Rows: 2 + +-- `on(host)` leaves only `host`, so `without(host)` groups everything together. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_a / on(host) outer_agg_b); + ++---------------------+--------------------------------------------------------------+ +| ts | sum(outer_agg_a.greptime_value / outer_agg_b.greptime_value) | ++---------------------+--------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+--------------------------------------------------------------+ + +-- `device` is not a label of the operand any more, so it groups everything together too. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_a / on(host) outer_agg_b); + ++---------------------+--------------------------------------------------------------+ +| ts | sum(outer_agg_a.greptime_value / outer_agg_b.greptime_value) | ++---------------------+--------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+--------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_a / on(host) outer_agg_b); + ++------+---------------------+--------------------------------------------------------------+ +| host | ts | sum(outer_agg_a.greptime_value / outer_agg_b.greptime_value) | ++------+---------------------+--------------------------------------------------------------+ +| h1 | 1970-01-01T00:00:00 | 5.0 | +| h2 | 1970-01-01T00:00:00 | 5.0 | ++------+---------------------+--------------------------------------------------------------+ + +DROP TABLE outer_agg_a; + +Affected Rows: 0 + +DROP TABLE outer_agg_b; + +Affected Rows: 0 + +-- Same over the metric engine, whose operands carry `__tsid`. +CREATE TABLE outer_agg_physical ( + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, +) ENGINE = metric WITH ("physical_metric_table" = ""); + +Affected Rows: 0 + +CREATE TABLE outer_agg_metric_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host, device), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +Affected Rows: 0 + +CREATE TABLE outer_agg_metric_b ( + host STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +Affected Rows: 0 + +INSERT INTO outer_agg_metric_a (host, device, ts, greptime_value) VALUES + ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +Affected Rows: 2 + +INSERT INTO outer_agg_metric_b (host, ts, greptime_value) VALUES + ('h1', 0, 2), ('h2', 0, 4); + +Affected Rows: 2 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + ++---------------------+----------------------------------------------------------------------------+ +| ts | sum(outer_agg_metric_a.greptime_value / outer_agg_metric_b.greptime_value) | ++---------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_metric_a / on(host) outer_agg_metric_b); + ++---------------------+----------------------------------------------------------------------------+ +| ts | sum(outer_agg_metric_a.greptime_value / outer_agg_metric_b.greptime_value) | ++---------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + ++------+---------------------+----------------------------------------------------------------------------+ +| host | ts | sum(outer_agg_metric_a.greptime_value / outer_agg_metric_b.greptime_value) | ++------+---------------------+----------------------------------------------------------------------------+ +| h1 | 1970-01-01T00:00:00 | 5.0 | +| h2 | 1970-01-01T00:00:00 | 5.0 | ++------+---------------------+----------------------------------------------------------------------------+ + +DROP TABLE outer_agg_metric_a; + +Affected Rows: 0 + +DROP TABLE outer_agg_metric_b; + +Affected Rows: 0 + +DROP TABLE outer_agg_physical; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/vector_matching.sql b/tests/cases/standalone/common/promql/vector_matching.sql new file mode 100644 index 00000000000..1fd0091c901 --- /dev/null +++ b/tests/cases/standalone/common/promql/vector_matching.sql @@ -0,0 +1,221 @@ +CREATE TABLE counter_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +CREATE TABLE gauge_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +INSERT INTO counter_metric VALUES + ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20), ('host2', 'eth0', 0, 30); + +INSERT INTO gauge_metric VALUES + ('host1', 'eth0', 0, 2), ('host1', 'eth1', 0, 4), ('host2', 'eth0', 0, 5); + +-- Matching on the whole tag set keeps it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / gauge_metric; + +-- `on(...)` reduces the result to the matching labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) gauge_metric{device="eth0"}; + +-- `ignoring(...)` drops the ignored labels, here from operands that don't even share them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / ignoring(device) gauge_metric{device="eth1"}; + +-- `on()` matches everything against everything, so both sides must hold a single series. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{host="host1",device="eth0"} / on() gauge_metric{host="host1",device="eth0"}; + +-- A filtering comparison keeps the left sample values and the derived labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > on(host) gauge_metric{device="eth0"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > bool on(host) gauge_metric{device="eth0"}; + +-- `host1` carries two devices on both sides: the match group repeats on the one side. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric; + +-- The one side is unique here, the many side is not, and no group modifier makes it explicit. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; + +-- `group_left` takes the result labels from the many side. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left gauge_metric{device="eth0"}; + +-- `group_right` swaps the sides. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_right gauge_metric; + +-- `group_right` makes the left operand the one side, which `host1` breaks. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_right gauge_metric{device="eth0"}; + +-- `group_left(device)` copies `device` from the one side, so both `host1` rows end up with the +-- same labels. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left(device) gauge_metric{device="eth0"}; + +-- An included label the one side doesn't carry is dropped from the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_left(device) sum by(host)(gauge_metric); + +DROP TABLE counter_metric; + +DROP TABLE gauge_metric; + +-- A label can carry the name the cardinality check generates for its row count. +CREATE TABLE collide_left ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +CREATE TABLE collide_right ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +INSERT INTO collide_left VALUES ('host1', 'a', 0, 10); + +INSERT INTO collide_right VALUES ('host1', 'b', 0, 2); + +TQL EVAL (0, 0, '5s') collide_left / on(host) collide_right; + +DROP TABLE collide_left; + +DROP TABLE collide_right; + +-- Partitioning on a column outside the match keys puts one match group in several regions. The +-- cardinality check counts the merged data, so it must still see the duplicate. +CREATE TABLE spread_left ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +) +PARTITION ON COLUMNS (device) ( + device < 'eth1', + device >= 'eth1' +); + +CREATE TABLE spread_right ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +INSERT INTO spread_left VALUES ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20); + +INSERT INTO spread_right VALUES ('host1', 0, 2); + +TQL EVAL (0, 0, '5s') spread_left / on(host) spread_right; + +TQL EVAL (0, 0, '5s') spread_right / on(host) group_left spread_left; + +DROP TABLE spread_left; + +DROP TABLE spread_right; + +-- An outer aggregate groups on the labels the operand has left, not on the ones the scan +-- underneath could still produce. +CREATE TABLE outer_agg_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +CREATE TABLE outer_agg_b ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +INSERT INTO outer_agg_a VALUES ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +INSERT INTO outer_agg_b VALUES ('h1', 0, 2), ('h2', 0, 4); + +-- `on(host)` leaves only `host`, so `without(host)` groups everything together. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_a / on(host) outer_agg_b); + +-- `device` is not a label of the operand any more, so it groups everything together too. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_a / on(host) outer_agg_b); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_a / on(host) outer_agg_b); + +DROP TABLE outer_agg_a; + +DROP TABLE outer_agg_b; + +-- Same over the metric engine, whose operands carry `__tsid`. +CREATE TABLE outer_agg_physical ( + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, +) ENGINE = metric WITH ("physical_metric_table" = ""); + +CREATE TABLE outer_agg_metric_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host, device), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +CREATE TABLE outer_agg_metric_b ( + host STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +INSERT INTO outer_agg_metric_a (host, device, ts, greptime_value) VALUES + ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +INSERT INTO outer_agg_metric_b (host, ts, greptime_value) VALUES + ('h1', 0, 2), ('h2', 0, 4); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_metric_a / on(host) outer_agg_metric_b); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + +DROP TABLE outer_agg_metric_a; + +DROP TABLE outer_agg_metric_b; + +DROP TABLE outer_agg_physical; diff --git a/tests/cases/standalone/common/skip_wal.result b/tests/cases/standalone/common/skip_wal.result index 25c18fed20e..792660b2f9a 100644 --- a/tests/cases/standalone/common/skip_wal.result +++ b/tests/cases/standalone/common/skip_wal.result @@ -19,50 +19,6 @@ VALUES Affected Rows: 3 --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - -+------+-----+----------+-------------+-----------+----+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+------+-----+----------+-------------+-----------+----+ -+------+-----+----------+-------------+-----------+----+ - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -ADMIN flush_table('system_metrics'); - -+-------------------------------------+ -| ADMIN flush_table('system_metrics') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - --- A table created without a real WAL provider cannot enable WAL later. -ALTER TABLE system_metrics SET 'skip_wal' = 'false'; - -Error: 1001(Unsupported), Unsupported operation setting skip_wal = 'false' without an existing WAL provider - -DROP TABLE system_metrics; - -Affected Rows: 0 - CREATE TABLE alter_skip_wal ( host STRING, val DOUBLE, @@ -152,8 +108,15 @@ SELECT * FROM alter_skip_wal ORDER BY ts; | host3 | 3.0 | 1970-01-01T00:00:03 | +-------+-----+---------------------+ --- Re-enabling WAL does not retroactively write the skipped row to WAL. -- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + ++------+-----+----------+-------------+-----------+----+ +| host | idc | cpu_util | memory_util | disk_util | ts | ++------+-----+----------+-------------+-----------+----+ ++------+-----+----------+-------------+-----------+----+ + +-- Re-enabling WAL does not retroactively write the skipped row to WAL. SHOW CREATE TABLE alter_skip_wal; +----------------+-----------------------------------------------+ @@ -186,3 +149,39 @@ DROP TABLE alter_skip_wal; Affected Rows: 0 +INSERT INTO system_metrics +VALUES + ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), + ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), + ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); + +Affected Rows: 3 + +ADMIN flush_table('system_metrics'); + ++-------------------------------------+ +| ADMIN flush_table('system_metrics') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +-- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + ++-------+-------+----------+-------------+-----------+-------------------------+ +| host | idc | cpu_util | memory_util | disk_util | ts | ++-------+-------+----------+-------------+-----------+-------------------------+ +| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | +| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | +| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | ++-------+-------+----------+-------------+-----------+-------------------------+ + +-- A table created without a real WAL provider cannot enable WAL later. +ALTER TABLE system_metrics SET 'skip_wal' = 'false'; + +Error: 1001(Unsupported), Unsupported operation setting skip_wal = 'false' without an existing WAL provider + +DROP TABLE system_metrics; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/skip_wal.sql b/tests/cases/standalone/common/skip_wal.sql index 2e9cac1f272..e0bfa4b41d8 100644 --- a/tests/cases/standalone/common/skip_wal.sql +++ b/tests/cases/standalone/common/skip_wal.sql @@ -15,25 +15,6 @@ VALUES ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -ADMIN flush_table('system_metrics'); - --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - --- A table created without a real WAL provider cannot enable WAL later. -ALTER TABLE system_metrics SET 'skip_wal' = 'false'; - -DROP TABLE system_metrics; - CREATE TABLE alter_skip_wal ( host STRING, val DOUBLE, @@ -65,10 +46,28 @@ INSERT INTO alter_skip_wal VALUES ('host3', 3, 3000); SELECT * FROM alter_skip_wal ORDER BY ts; --- Re-enabling WAL does not retroactively write the skipped row to WAL. -- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + +-- Re-enabling WAL does not retroactively write the skipped row to WAL. SHOW CREATE TABLE alter_skip_wal; SELECT * FROM alter_skip_wal ORDER BY ts; DROP TABLE alter_skip_wal; + +INSERT INTO system_metrics +VALUES + ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), + ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), + ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); + +ADMIN flush_table('system_metrics'); + +-- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + +-- A table created without a real WAL provider cannot enable WAL later. +ALTER TABLE system_metrics SET 'skip_wal' = 'false'; + +DROP TABLE system_metrics; diff --git a/tests/cases/standalone/common/tql/matching_filter.result b/tests/cases/standalone/common/tql/matching_filter.result index 8f5bebbe8c3..c0df22a738e 100644 --- a/tests/cases/standalone/common/tql/matching_filter.result +++ b/tests/cases/standalone/common/tql/matching_filter.result @@ -179,24 +179,11 @@ tql eval(0, 5, '5s') counter_metric offset 5s / on(host, device) gauge_metric{ho | host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | +-------+--------+---------------------+---------------------------------------+ --- Matching on a subset of the labels. Prometheus rejects this as many-to-many; GreptimeDB --- returns the cross product instead (#9209). Recorded here to show the rewrite does not --- change it, not to endorse it. +-- Matching on a subset of the labels: both operands hold several devices per host. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{host="host1"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side -- `ignoring(...)`: every label that is not ignored is a matching label. -- SQLNESS SORT_RESULT 3 1 @@ -214,48 +201,19 @@ tql eval(0, 5, '5s') counter_metric / ignoring(missing_label) gauge_metric{host= -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{host="host1"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side -- An ignored label is not a matching label: the left operand keeps its `eth1` series. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{device="eth0"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | -| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) -- Same for a label left out of `on(...)`. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | -| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) -- Matcher kinds other than equality are enforced by the join just the same. -- SQLNESS SORT_RESULT 3 1 @@ -321,16 +279,7 @@ tql eval(0, 5, '5s') sum by(host)(counter_metric{device="eth0"}) / on(host) sum -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) sum by(host)(gauge_metric{device="eth0"}); -+-------+---------------------+---------------------------------------------------------+ -| host | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | -+-------+---------------------+---------------------------------------------------------+ -| host1 | 1970-01-01T00:00:00 | 10.0 | -| host1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | 1970-01-01T00:00:05 | 10.0 | -| host1 | 1970-01-01T00:00:05 | 5.0 | -| host2 | 1970-01-01T00:00:00 | 20.0 | -| host2 | 1970-01-01T00:00:05 | 20.0 | -+-------+---------------------+---------------------------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) -- A grouping label that is a value field, not a tag: its value varies between samples of one -- series, so a matcher on it must not filter the scan before sample selection (#9242). @@ -411,14 +360,14 @@ tql eval(0, 5, '5s') counter_metric > on(host, device) gauge_metric{host="host1" -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) group_left sum by(host)(gauge_metric{host="host1"}); -+-------+---------------------+---------------------------------------------------------+ -| host | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | -+-------+---------------------+---------------------------------------------------------+ -| host1 | 1970-01-01T00:00:00 | 2.5 | -| host1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | 1970-01-01T00:00:05 | 2.5 | -| host1 | 1970-01-01T00:00:05 | 5.0 | -+-------+---------------------+---------------------------------------------------------+ ++-------+--------+---------------------+---------------------------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | ++-------+--------+---------------------+---------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 2.5 | +| host1 | eth0 | 1970-01-01T00:00:05 | 2.5 | +| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | ++-------+--------+---------------------+---------------------------------------------------------+ -- `absent()` reports the inner selector's equality matchers as labels, so a copied matcher -- must not reach the context an enclosing expression reads. diff --git a/tests/cases/standalone/common/tql/matching_filter.sql b/tests/cases/standalone/common/tql/matching_filter.sql index 72da3da70a0..1b8492c4351 100644 --- a/tests/cases/standalone/common/tql/matching_filter.sql +++ b/tests/cases/standalone/common/tql/matching_filter.sql @@ -73,9 +73,7 @@ tql eval(0, 5, '5s') count_over_time(counter_metric[10s]) / on(host, device) cou -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric offset 5s / on(host, device) gauge_metric{host="host1"}; --- Matching on a subset of the labels. Prometheus rejects this as many-to-many; GreptimeDB --- returns the cross product instead (#9209). Recorded here to show the rewrite does not --- change it, not to endorse it. +-- Matching on a subset of the labels: both operands hold several devices per host. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{host="host1"}; diff --git a/tests/cases/standalone/common/ttl/ttl_instant.result b/tests/cases/standalone/common/ttl/ttl_instant.result index 110fd92d92b..d6a3921d245 100644 --- a/tests/cases/standalone/common/ttl/ttl_instant.result +++ b/tests/cases/standalone/common/ttl/ttl_instant.result @@ -45,7 +45,6 @@ ORDER BY +-----+ +-----+ --- SQLNESS SLEEP 2s ADMIN flush_table('test_ttl'); +-------------------------------+ diff --git a/tests/cases/standalone/common/ttl/ttl_instant.sql b/tests/cases/standalone/common/ttl/ttl_instant.sql index 57bc4b625ea..be4831ba2b4 100644 --- a/tests/cases/standalone/common/ttl/ttl_instant.sql +++ b/tests/cases/standalone/common/ttl/ttl_instant.sql @@ -20,7 +20,6 @@ from ORDER BY val; --- SQLNESS SLEEP 2s ADMIN flush_table('test_ttl'); ADMIN compact_table('test_ttl'); diff --git a/tests/cases/standalone/common/types/json/json2.result b/tests/cases/standalone/common/types/json/json2.result index ce0c9748ce1..7114d9bafa1 100644 --- a/tests/cases/standalone/common/types/json/json2.result +++ b/tests/cases/standalone/common/types/json/json2.result @@ -513,3 +513,65 @@ drop table json2_finite_paths; Affected Rows: 0 +-- A newly added JSON2 column must not read data from a dropped column with the +-- same name in an older SST. Column identity is determined by column id, not +-- by the parquet root name. +create table json2_drop_add_same_name ( + ts timestamp time index, + b json2 +) with ( + 'append_mode' = 'true', + 'sst_format' = 'flat', +); + +Affected Rows: 0 + +insert into json2_drop_add_same_name +values (1, '{"a": "old"}'); + +Affected Rows: 1 + +admin flush_table('json2_drop_add_same_name'); + ++-----------------------------------------------+ +| ADMIN flush_table('json2_drop_add_same_name') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +alter table json2_drop_add_same_name drop column b; + +Affected Rows: 0 + +alter table json2_drop_add_same_name add column b json2; + +Affected Rows: 0 + +insert into json2_drop_add_same_name +values (2, '{"a": "new"}'); + +Affected Rows: 1 + +admin flush_table('json2_drop_add_same_name'); + ++-----------------------------------------------+ +| ADMIN flush_table('json2_drop_add_same_name') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +select ts, b.a +from json2_drop_add_same_name +order by ts; + ++-------------------------+--------------------------------------------------+ +| ts | json_get(json2_drop_add_same_name.b,Utf8("$.a")) | ++-------------------------+--------------------------------------------------+ +| 1970-01-01T00:00:00.001 | | +| 1970-01-01T00:00:00.002 | new | ++-------------------------+--------------------------------------------------+ + +drop table json2_drop_add_same_name; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2.sql b/tests/cases/standalone/common/types/json/json2.sql index 17248741c35..55642b7530c 100644 --- a/tests/cases/standalone/common/types/json/json2.sql +++ b/tests/cases/standalone/common/types/json/json2.sql @@ -157,3 +157,33 @@ from json2_finite_paths order by ts; drop table json2_finite_paths; + +-- A newly added JSON2 column must not read data from a dropped column with the +-- same name in an older SST. Column identity is determined by column id, not +-- by the parquet root name. +create table json2_drop_add_same_name ( + ts timestamp time index, + b json2 +) with ( + 'append_mode' = 'true', + 'sst_format' = 'flat', +); + +insert into json2_drop_add_same_name +values (1, '{"a": "old"}'); + +admin flush_table('json2_drop_add_same_name'); + +alter table json2_drop_add_same_name drop column b; +alter table json2_drop_add_same_name add column b json2; + +insert into json2_drop_add_same_name +values (2, '{"a": "new"}'); + +admin flush_table('json2_drop_add_same_name'); + +select ts, b.a +from json2_drop_add_same_name +order by ts; + +drop table json2_drop_add_same_name; diff --git a/tests/cases/standalone/common/types/json/json2_alter.result b/tests/cases/standalone/common/types/json/json2_alter.result index c23f9e2533f..055ea90346a 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.result +++ b/tests/cases/standalone/common/types/json/json2_alter.result @@ -2,37 +2,110 @@ CREATE TABLE application_logs ( ts TIMESTAMP TIME INDEX, attrs JSON2 ) WITH ( - 'append_mode' = 'true' + 'append_mode' = 'true', + 'sst_format' = 'flat' ); Affected Rows: 0 -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2; +INSERT INTO application_logs VALUES + (1, '{"service":"frontend","duration":12,"trace_id":"before-1"}'), + (2, '{"service":"ingester","duration":34,"trace_id":"before-2"}'); -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 +Affected Rows: 2 ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( - max_auto_expanded_paths = 2000 + max_auto_expanded_paths = 2000, + service STRING, + duration BIGINT, + trace_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +SELECT COUNT(*) AS alter_flushed_sst_num +FROM information_schema.ssts_manifest +WHERE visible + AND table_id IN ( + SELECT table_id FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = 'application_logs' + ); + ++-----------------------+ +| alter_flushed_sst_num | ++-----------------------+ +| 1 | ++-----------------------+ + +INSERT INTO application_logs VALUES + (3, '{"service":"frontend","duration":56,"trace_id":"after-1"}'), + (4, '{"service":"compactor","trace_id":"after-2"}'); + +Affected Rows: 2 + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.service"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.duration"),Int64(NULL)) | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | frontend | 12 | before-1 | +| 1970-01-01T00:00:00.002 | ingester | 34 | before-2 | +| 1970-01-01T00:00:00.003 | frontend | 56 | after-1 | +| 1970-01-01T00:00:00.004 | compactor | | after-2 | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ + +ADMIN FLUSH_TABLE('application_logs'); + ++---------------------------------------+ +| ADMIN FLUSH_TABLE('application_logs') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.service"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.duration"),Int64(NULL)) | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | frontend | 12 | before-1 | +| 1970-01-01T00:00:00.002 | ingester | 34 | before-2 | +| 1970-01-01T00:00:00.003 | frontend | 56 | after-1 | +| 1970-01-01T00:00:00.004 | compactor | | after-2 | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +INSERT INTO application_logs VALUES + (5, '{"trace_id":"after-3","user":{"id":"u1","name":"alice"},"request_id":"r1"}'); + +Affected Rows: 1 + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.name"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.request_id"),Utf8View(NULL)) | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | before-1 | | | | +| 1970-01-01T00:00:00.002 | before-2 | | | | +| 1970-01-01T00:00:00.003 | after-1 | | | | +| 1970-01-01T00:00:00.004 | after-2 | | | | +| 1970-01-01T00:00:00.005 | after-3 | u1 | alice | r1 | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( @@ -40,12 +113,127 @@ ALTER TABLE application_logs trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +INSERT INTO application_logs VALUES + (6, '{"trace_id":"after-4","user":{"id":"u2"},"request_id":"r2"}'); + +Affected Rows: 1 + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.name"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.request_id"),Utf8View(NULL)) | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | before-1 | | | | +| 1970-01-01T00:00:00.002 | before-2 | | | | +| 1970-01-01T00:00:00.003 | after-1 | | | | +| 1970-01-01T00:00:00.004 | after-2 | | | | +| 1970-01-01T00:00:00.005 | after-3 | u1 | alice | r1 | +| 1970-01-01T00:00:00.006 | after-4 | u2 | | r2 | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ + +ADMIN FLUSH_TABLE('application_logs'); + ++---------------------------------------+ +| ADMIN FLUSH_TABLE('application_logs') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ DROP TABLE application_logs; Affected Rows: 0 +create table json2_alter_add_settings ( + ts timestamp time index +) with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +-- ADD COLUMN preserves JSON2 settings and type hints, including for existing rows. +insert into json2_alter_add_settings values (1); + +Affected Rows: 1 + +alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); + +Affected Rows: 0 + +alter table json2_alter_add_settings add column k json2(service string); + +Affected Rows: 0 + +alter table json2_alter_add_settings add column m json2(max_auto_expanded_paths = 0, nested.value bigint); + +Affected Rows: 0 + +show create table json2_alter_add_settings; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| json2_alter_add_settings | CREATE TABLE IF NOT EXISTS "json2_alter_add_settings" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100 | +| | ) NULL, | +| | "k" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "service" STRING | +| | ) NULL, | +| | "m" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "nested"."value" BIGINT | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +insert into json2_alter_add_settings values + (2, '{"a":1}', '{"service":"api"}', '{"nested":{"value":42},"extra":true}'); + +Affected Rows: 1 + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| ts | j | json_get(json2_alter_add_settings.k,Utf8("$.service"),Utf8View(NULL)) | json_get(json2_alter_add_settings.m,Utf8("$.nested.value"),Int64(NULL)) | m | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| 1970-01-01T00:00:00.001 | | | | | +| 1970-01-01T00:00:00.002 | {"a":1} | api | 42 | {"extra":true,"nested":{"value":42}} | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ + +admin flush_table('json2_alter_add_settings'); + ++-----------------------------------------------+ +| ADMIN flush_table('json2_alter_add_settings') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| ts | j | json_get(json2_alter_add_settings.k,Utf8("$.service"),Utf8View(NULL)) | json_get(json2_alter_add_settings.m,Utf8("$.nested.value"),Int64(NULL)) | m | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| 1970-01-01T00:00:00.001 | | | | | +| 1970-01-01T00:00:00.002 | {"a":1} | api | 42 | {"extra":true,"nested":{"value":42}} | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ + +drop table json2_alter_add_settings; + +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 index c89556c3485..83b514f6895 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.sql +++ b/tests/cases/standalone/common/types/json/json2_alter.sql @@ -2,35 +2,102 @@ CREATE TABLE application_logs ( ts TIMESTAMP TIME INDEX, attrs JSON2 ) WITH ( - 'append_mode' = 'true' + 'append_mode' = 'true', + 'sst_format' = 'flat' ); -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2; - -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2 (); +INSERT INTO application_logs VALUES + (1, '{"service":"frontend","duration":12,"trace_id":"before-1"}'), + (2, '{"service":"ingester","duration":34,"trace_id":"before-2"}'); ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( - max_auto_expanded_paths = 2000 + max_auto_expanded_paths = 2000, + service STRING, + duration BIGINT, + trace_id STRING ); +SELECT COUNT(*) AS alter_flushed_sst_num +FROM information_schema.ssts_manifest +WHERE visible + AND table_id IN ( + SELECT table_id FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = 'application_logs' + ); + +INSERT INTO application_logs VALUES + (3, '{"service":"frontend","duration":56,"trace_id":"after-1"}'), + (4, '{"service":"compactor","trace_id":"after-2"}'); + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + +ADMIN FLUSH_TABLE('application_logs'); + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); +INSERT INTO application_logs VALUES + (5, '{"trace_id":"after-3","user":{"id":"u1","name":"alice"},"request_id":"r1"}'); + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); +INSERT INTO application_logs VALUES + (6, '{"trace_id":"after-4","user":{"id":"u2"},"request_id":"r2"}'); + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + +ADMIN FLUSH_TABLE('application_logs'); + DROP TABLE application_logs; + +create table json2_alter_add_settings ( + ts timestamp time index +) with ( + 'append_mode' = 'true' +); + +-- ADD COLUMN preserves JSON2 settings and type hints, including for existing rows. +insert into json2_alter_add_settings values (1); +alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); + +alter table json2_alter_add_settings add column k json2(service string); + +alter table json2_alter_add_settings add column m json2(max_auto_expanded_paths = 0, nested.value bigint); + +show create table json2_alter_add_settings; + +insert into json2_alter_add_settings values + (2, '{"a":1}', '{"service":"api"}', '{"nested":{"value":42},"extra":true}'); + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + +admin flush_table('json2_alter_add_settings'); + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + +drop table json2_alter_add_settings; diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result new file mode 100644 index 00000000000..aaefd5ff0ca --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result @@ -0,0 +1,83 @@ +CREATE TABLE json2_alter_type_hints_compaction ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction VALUES + (1, '{"kind":1,"message":"before-has"}'), + (2, '{"message":"before-missing"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + ++--------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +ALTER TABLE json2_alter_type_hints_compaction + MODIFY COLUMN attrs JSON2 ( + kind BIGINT + ); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction VALUES + (3, '{"kind":3,"message":"after-has"}'), + (4, '{"message":"after-missing"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + ++--------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.message")) | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | before-has | +| 1970-01-01T00:00:00.002 | | before-missing | +| 1970-01-01T00:00:00.003 | 3 | after-has | +| 1970-01-01T00:00:00.004 | | after-missing | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400'); + ++---------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400') | ++---------------------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.message")) | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | before-has | +| 1970-01-01T00:00:00.002 | | before-missing | +| 1970-01-01T00:00:00.003 | 3 | after-has | +| 1970-01-01T00:00:00.004 | | after-missing | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ + +DROP TABLE json2_alter_type_hints_compaction; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql new file mode 100644 index 00000000000..34a16dafbb7 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql @@ -0,0 +1,36 @@ +CREATE TABLE json2_alter_type_hints_compaction ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_alter_type_hints_compaction VALUES + (1, '{"kind":1,"message":"before-has"}'), + (2, '{"message":"before-missing"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + +ALTER TABLE json2_alter_type_hints_compaction + MODIFY COLUMN attrs JSON2 ( + kind BIGINT + ); + +INSERT INTO json2_alter_type_hints_compaction VALUES + (3, '{"kind":3,"message":"after-has"}'), + (4, '{"message":"after-missing"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + +DROP TABLE json2_alter_type_hints_compaction; diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result new file mode 100644 index 00000000000..e3517c90e38 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result @@ -0,0 +1,78 @@ +CREATE TABLE json2_alter_type_hints_compaction_conversion ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (1, '{"to_string":123,"to_int":"456"}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + ++-------------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion') | ++-------------------------------------------------------------------+ +| 0 | ++-------------------------------------------------------------------+ + +ALTER TABLE json2_alter_type_hints_compaction_conversion + MODIFY COLUMN attrs JSON2 ( + to_string STRING, + to_int BIGINT + ); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (2, '{"to_string":"after","to_int":789}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + ++-------------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion') | ++-------------------------------------------------------------------+ +| 0 | ++-------------------------------------------------------------------+ + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_string"),Utf8View(NULL)) | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_int"),Int64(NULL)) | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 123 | 456 | +| 1970-01-01T00:00:00.002 | after | 789 | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400'); + ++--------------------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400') | ++--------------------------------------------------------------------------------------+ +| 0 | ++--------------------------------------------------------------------------------------+ + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_string"),Utf8View(NULL)) | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_int"),Int64(NULL)) | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 123 | 456 | +| 1970-01-01T00:00:00.002 | after | 789 | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ + +DROP TABLE json2_alter_type_hints_compaction_conversion; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql new file mode 100644 index 00000000000..cfadce89a93 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql @@ -0,0 +1,35 @@ +CREATE TABLE json2_alter_type_hints_compaction_conversion ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (1, '{"to_string":123,"to_int":"456"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + +ALTER TABLE json2_alter_type_hints_compaction_conversion + MODIFY COLUMN attrs JSON2 ( + to_string STRING, + to_int BIGINT + ); + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (2, '{"to_string":"after","to_int":789}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400'); + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + +DROP TABLE json2_alter_type_hints_compaction_conversion; diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result new file mode 100644 index 00000000000..5ad7a46ce83 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result @@ -0,0 +1,80 @@ +CREATE TABLE json2_compaction_null_type_hint_mismatch ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (1, '{"kind":1}'), + (2, '{"kind":"invalid","message":"keep"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + ++---------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch') | ++---------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------+ + +ALTER TABLE json2_compaction_null_type_hint_mismatch + MODIFY COLUMN attrs JSON2 ( + kind BIGINT + ); + +Affected Rows: 0 + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (3, '{"kind":3}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + ++---------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch') | ++---------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| ts | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.message")) | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | | +| 1970-01-01T00:00:00.002 | | keep | +| 1970-01-01T00:00:00.003 | 3 | | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400'); + ++----------------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400') | ++----------------------------------------------------------------------------------+ +| 0 | ++----------------------------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| ts | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.message")) | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | | +| 1970-01-01T00:00:00.002 | | keep | +| 1970-01-01T00:00:00.003 | 3 | | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ + +DROP TABLE json2_compaction_null_type_hint_mismatch; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql new file mode 100644 index 00000000000..8025977c380 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql @@ -0,0 +1,35 @@ +CREATE TABLE json2_compaction_null_type_hint_mismatch ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (1, '{"kind":1}'), + (2, '{"kind":"invalid","message":"keep"}'); + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + +ALTER TABLE json2_compaction_null_type_hint_mismatch + MODIFY COLUMN attrs JSON2 ( + kind BIGINT + ); + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (3, '{"kind":3}'); + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + +DROP TABLE json2_compaction_null_type_hint_mismatch; diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.result b/tests/cases/standalone/common/types/json/json2_type_hints.result index 1a892b6e4cb..eda5aee7692 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.result +++ b/tests/cases/standalone/common/types/json/json2_type_hints.result @@ -71,7 +71,7 @@ CREATE TABLE json2_type_hints_timestamp ( ) ); -Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TIMESTAMP +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TIMESTAMP; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 -- A type hint at the maximum supported depth is accepted. CREATE TABLE json2_type_hint_depth_50 ( @@ -105,3 +105,132 @@ DROP TABLE json2_type_hint_depth_50; Affected Rows: 0 +-- Type hints reject unsupported types and aliases. +CREATE TABLE json2_hint_int8 (ts TIMESTAMP TIME INDEX, j JSON2 (value INT8)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: INT8; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_tinyint (ts TIMESTAMP TIME INDEX, j JSON2 (value TINYINT)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TINYINT; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_uint8 (ts TIMESTAMP TIME INDEX, j JSON2 (value UINT8)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: UInt8; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_float (ts TIMESTAMP TIME INDEX, j JSON2 (value FLOAT)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: FLOAT; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_text (ts TIMESTAMP TIME INDEX, j JSON2 (value TEXT)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TEXT; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_unsigned ( + ts TIMESTAMP TIME INDEX, + j JSON2 (value BIGINT UNSIGNED) +) WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +INSERT INTO json2_hint_unsigned VALUES (1, '{"value":18446744073709551615}'); + +Affected Rows: 1 + +SELECT j.value FROM json2_hint_unsigned; + ++--------------------------------------------------------------+ +| json_get(json2_hint_unsigned.j,Utf8("$.value"),UInt64(NULL)) | ++--------------------------------------------------------------+ +| 18446744073709551615 | ++--------------------------------------------------------------+ + +ALTER TABLE json2_hint_unsigned MODIFY COLUMN j JSON2 (value INT8); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: INT8; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +DROP TABLE json2_hint_unsigned; + +Affected Rows: 0 + +-- These aliases preserve the exact type and are displayed using canonical names. +CREATE TABLE json2_hint_aliases ( + ts TIMESTAMP TIME INDEX, + j JSON2 (i INT64, u UINT64, f FLOAT64) +) WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_hint_aliases; + ++--------------------+---------------------------------------------------+ +| Table | Create Table | ++--------------------+---------------------------------------------------+ +| json2_hint_aliases | CREATE TABLE IF NOT EXISTS "json2_hint_aliases" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "i" BIGINT, | +| | "u" BIGINT UNSIGNED, | +| | "f" DOUBLE | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------+---------------------------------------------------+ + +INSERT INTO json2_hint_aliases VALUES + (1, '{"i":-9223372036854775808,"u":18446744073709551615,"f":1.5}'); + +Affected Rows: 1 + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + ++----------------------+----------------------+-----+ +| i | u | f | ++----------------------+----------------------+-----+ +| -9223372036854775808 | 18446744073709551615 | 1.5 | ++----------------------+----------------------+-----+ + +ALTER TABLE json2_hint_aliases MODIFY COLUMN j JSON2 (i INT64, u UINT64, f FLOAT64); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_hint_aliases; + ++--------------------+---------------------------------------------------+ +| Table | Create Table | ++--------------------+---------------------------------------------------+ +| json2_hint_aliases | CREATE TABLE IF NOT EXISTS "json2_hint_aliases" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "i" BIGINT, | +| | "u" BIGINT UNSIGNED, | +| | "f" DOUBLE | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------+---------------------------------------------------+ + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + ++----------------------+----------------------+-----+ +| i | u | f | ++----------------------+----------------------+-----+ +| -9223372036854775808 | 18446744073709551615 | 1.5 | ++----------------------+----------------------+-----+ + +DROP TABLE json2_hint_aliases; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.sql b/tests/cases/standalone/common/types/json/json2_type_hints.sql index aa88ae60bf0..3b301700333 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.sql +++ b/tests/cases/standalone/common/types/json/json2_type_hints.sql @@ -59,3 +59,48 @@ CREATE TABLE json2_type_hint_depth_51 ( DROP TABLE json2_type_hints; DROP TABLE json2_type_hint_depth_50; + +-- Type hints reject unsupported types and aliases. +CREATE TABLE json2_hint_int8 (ts TIMESTAMP TIME INDEX, j JSON2 (value INT8)); + +CREATE TABLE json2_hint_tinyint (ts TIMESTAMP TIME INDEX, j JSON2 (value TINYINT)); + +CREATE TABLE json2_hint_uint8 (ts TIMESTAMP TIME INDEX, j JSON2 (value UINT8)); + +CREATE TABLE json2_hint_float (ts TIMESTAMP TIME INDEX, j JSON2 (value FLOAT)); + +CREATE TABLE json2_hint_text (ts TIMESTAMP TIME INDEX, j JSON2 (value TEXT)); + +CREATE TABLE json2_hint_unsigned ( + ts TIMESTAMP TIME INDEX, + j JSON2 (value BIGINT UNSIGNED) +) WITH ('append_mode' = 'true'); + +INSERT INTO json2_hint_unsigned VALUES (1, '{"value":18446744073709551615}'); + +SELECT j.value FROM json2_hint_unsigned; + +ALTER TABLE json2_hint_unsigned MODIFY COLUMN j JSON2 (value INT8); + +DROP TABLE json2_hint_unsigned; + +-- These aliases preserve the exact type and are displayed using canonical names. +CREATE TABLE json2_hint_aliases ( + ts TIMESTAMP TIME INDEX, + j JSON2 (i INT64, u UINT64, f FLOAT64) +) WITH ('append_mode' = 'true'); + +SHOW CREATE TABLE json2_hint_aliases; + +INSERT INTO json2_hint_aliases VALUES + (1, '{"i":-9223372036854775808,"u":18446744073709551615,"f":1.5}'); + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + +ALTER TABLE json2_hint_aliases MODIFY COLUMN j JSON2 (i INT64, u UINT64, f FLOAT64); + +SHOW CREATE TABLE json2_hint_aliases; + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + +DROP TABLE json2_hint_aliases; diff --git a/tests/cases/standalone/common/types/vector/vector.result b/tests/cases/standalone/common/types/vector/vector.result index bc9a38c8167..1da7ad913ac 100644 --- a/tests/cases/standalone/common/types/vector/vector.result +++ b/tests/cases/standalone/common/types/vector/vector.result @@ -18,6 +18,34 @@ INSERT INTO t VALUES Affected Rows: 3 +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.004 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + +ADMIN FLUSH_TABLE('t'); + ++------------------------+ +| ADMIN FLUSH_TABLE('t') | ++------------------------+ +| 0 | ++------------------------+ + +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.004 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + SELECT ts, v, vec_to_string(v) FROM t; +-------------------------+--------------------------+--------------------+ diff --git a/tests/cases/standalone/common/types/vector/vector.sql b/tests/cases/standalone/common/types/vector/vector.sql index 1f81d5cd56f..46af4661947 100644 --- a/tests/cases/standalone/common/types/vector/vector.sql +++ b/tests/cases/standalone/common/types/vector/vector.sql @@ -12,6 +12,12 @@ INSERT INTO t VALUES (5, parse_vec('[4.0, 5.0, 6.0]')), (6, parse_vec('[7.0, 8.0, 9.0]')); +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + +ADMIN FLUSH_TABLE('t'); + +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + SELECT ts, v, vec_to_string(v) FROM t; SELECT round(vec_cos_distance(v, '[0.0, 0.0, 0.0]'), 2) FROM t; diff --git a/tests/cases/standalone/flow-tql/flow_tql.result b/tests/cases/standalone/flow-tql/flow_tql.result index de93fb7c518..73321d040e1 100644 --- a/tests/cases/standalone/flow-tql/flow_tql.result +++ b/tests/cases/standalone/flow-tql/flow_tql.result @@ -155,7 +155,7 @@ Error: 2000(InvalidSyntax), Invalid TQL syntax: sql parser error: Expected expre CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-(now()+'15s'::interval), '5s') count_values("status_code", http_requests); -Error: 3001(EngineExecuteQuery), Failed to convert float seconds to duration, raw: -15: cannot convert float seconds to Duration: value is negative +Error: 3001(EngineExecuteQuery), The end time must be greater than start time, start: 0, end: -15000 CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-now()+'15s'::interval, '5s') count_values("status_code", http_requests); @@ -228,14 +228,14 @@ SELECT * FROM cnt_reqs ORDER BY ts, status_code; +--------------------------+---------------------+-------------+ | count(http_requests.val) | ts | status_code | +--------------------------+---------------------+-------------+ -| 3.0 | 1970-01-01T00:00:00 | 200.0 | -| 1.0 | 1970-01-01T00:00:00 | 401.0 | -| 1.0 | 1970-01-01T00:00:05 | 401.0 | -| 2.0 | 1970-01-01T00:00:05 | 404.0 | -| 1.0 | 1970-01-01T00:00:05 | 500.0 | -| 2.0 | 1970-01-01T00:00:10 | 200.0 | -| 2.0 | 1970-01-01T00:00:10 | 201.0 | -| 4.0 | 1970-01-01T00:00:15 | 500.0 | +| 3.0 | 1970-01-01T00:00:00 | 200 | +| 1.0 | 1970-01-01T00:00:00 | 401 | +| 1.0 | 1970-01-01T00:00:05 | 401 | +| 2.0 | 1970-01-01T00:00:05 | 404 | +| 1.0 | 1970-01-01T00:00:05 | 500 | +| 2.0 | 1970-01-01T00:00:10 | 200 | +| 2.0 | 1970-01-01T00:00:10 | 201 | +| 4.0 | 1970-01-01T00:00:15 | 500 | +--------------------------+---------------------+-------------+ DROP FLOW calc_reqs; diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.result b/tests/cases/standalone/tql-explain-analyze/tsid_column.result index c761267bd5a..3b509cdd20b 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.result +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.result @@ -54,7 +54,7 @@ TQL ANALYZE (0, 10, '5s') sum(tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -82,7 +82,7 @@ TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, instance@3 as instance, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -120,7 +120,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -133,7 +133,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -172,7 +172,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -185,7 +185,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/compat/README.md b/tests/compat/README.md deleted file mode 100644 index 07c4cb494f1..00000000000 --- a/tests/compat/README.md +++ /dev/null @@ -1,31 +0,0 @@ -# GreptimeDB compatibility test - -The compatibility test check whether a newer version of GreptimeDB can read from the data written by an old version of -GreptimeDB (the "backward" compatibility), and vice-versa (the "forward" compatibility). It's often ran in the Github -Actions to ensure there are no breaking changes by accident. - -The test work like this: reuse the sqlness-runner two times but each for a read or write side. For example, if we are -testing backward compatibility, we use sqlness-runner to run the SQLs with writes against the old GreptimeDB binary, and -use the same sqlness-runner to run the SQLs with reads against the new GreptimeDB binary. If the reads were executed -expectedly, we have achieved backward compatibility. - -This compatibility test is inspired by [Databend](https://github.com/datafuselabs/databend/). - -## Usage - -```shell -tests/compat/test-compat.sh -``` - -E.g. `tests/compat/test-compat.sh 0.6.0` tests if the data written by GreptimeDB **v0.6.0** can be read by **current** -version of GreptimeDB, and vice-versa. By "current", it's meant the fresh binary built by current codes. - -## Prerequisites - -Current version of GreptimeDB's binaries must reside in `./bins`: - -- `./bins/current/greptime` -- `./bins/current/sqlness-runner` - -The steps in Github Action already assure that. When running in local host, you have to `cp` them from target directory -manually. diff --git a/tests/compat/case/read/standalone/read.result b/tests/compat/case/read/standalone/read.result deleted file mode 100644 index ba59e7a4628..00000000000 --- a/tests/compat/case/read/standalone/read.result +++ /dev/null @@ -1,13 +0,0 @@ -select ts, i, s, f from foo order by ts; - -+---------------------+---+----------+-----+ -| ts | i | s | f | -+---------------------+---+----------+-----+ -| 2024-02-01T17:00:00 | 1 | my_tag_1 | | -| 2024-02-01T18:00:00 | 2 | my_tag_2 | | -| 2024-02-01T19:00:00 | 3 | my_tag_3 | | -| 2024-02-01T20:00:00 | 4 | my_tag_4 | 4.4 | -| 2024-02-01T21:00:00 | 5 | my_tag_5 | 5.5 | -| 2024-02-01T22:00:00 | 6 | my_tag_6 | 6.6 | -+---------------------+---+----------+-----+ - diff --git a/tests/compat/case/read/standalone/read.sql b/tests/compat/case/read/standalone/read.sql deleted file mode 100644 index e83d62d1ce9..00000000000 --- a/tests/compat/case/read/standalone/read.sql +++ /dev/null @@ -1 +0,0 @@ -select ts, i, s, f from foo order by ts; diff --git a/tests/compat/case/write/standalone/write.result b/tests/compat/case/write/standalone/write.result deleted file mode 100644 index 85ea219981c..00000000000 --- a/tests/compat/case/write/standalone/write.result +++ /dev/null @@ -1,25 +0,0 @@ -create table foo(ts timestamp time index, s string primary key, i int); - -Affected Rows: 0 - -insert into foo values -("2024-02-02 01:00:00+0800", "my_tag_1", 1), -("2024-02-02 02:00:00+0800", "my_tag_2", 2), -("2024-02-02 03:00:00+0800", "my_tag_3", 3); - -Affected Rows: 3 - --- Alter the table to trigger a flush (will be executed before process being terminated). --- Otherwise the SST might not be generated (the data could be remained in WAL). --- If we have the explicitly flush table interface in the future, it's still good to have the alter table in the test. -alter table foo add column f float; - -Affected Rows: 0 - -insert into foo values -("2024-02-02 04:00:00+0800", "my_tag_4", 4, 4.4), -("2024-02-02 05:00:00+0800", "my_tag_5", 5, 5.5), -("2024-02-02 06:00:00+0800", "my_tag_6", 6, 6.6); - -Affected Rows: 3 - diff --git a/tests/compat/case/write/standalone/write.sql b/tests/compat/case/write/standalone/write.sql deleted file mode 100644 index d7403a8e267..00000000000 --- a/tests/compat/case/write/standalone/write.sql +++ /dev/null @@ -1,16 +0,0 @@ -create table foo(ts timestamp time index, s string primary key, i int); - -insert into foo values -("2024-02-02 01:00:00+0800", "my_tag_1", 1), -("2024-02-02 02:00:00+0800", "my_tag_2", 2), -("2024-02-02 03:00:00+0800", "my_tag_3", 3); - --- Alter the table to trigger a flush (will be executed before process being terminated). --- Otherwise the SST might not be generated (the data could be remained in WAL). --- If we have the explicitly flush table interface in the future, it's still good to have the alter table in the test. -alter table foo add column f float; - -insert into foo values -("2024-02-02 04:00:00+0800", "my_tag_4", 4, 4.4), -("2024-02-02 05:00:00+0800", "my_tag_5", 5, 5.5), -("2024-02-02 06:00:00+0800", "my_tag_6", 6, 6.6); diff --git a/tests/compat/test-compat.sh b/tests/compat/test-compat.sh deleted file mode 100755 index 837a523b2ec..00000000000 --- a/tests/compat/test-compat.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/bin/bash - -set -o errexit - -usage() { - echo " Tests the compatibility between different versions of GreptimeDB." - echo " Expects the directory './bins/current' contains the newly built binaries." - echo " Usage: $0 " -} - -# The previous version of GreptimeDB to test compatibility with. -# e.g. old_ver="0.6.0" -old_ver="$1" - -if [ -z $old_ver ] -then - usage - exit -1 -fi - -SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)" -echo " === SCRIPT_PATH: $SCRIPT_PATH" -source "${SCRIPT_PATH}/util.sh" - -# go to work tree root -cd "$SCRIPT_PATH/../../" - -download_binary "$old_ver" - -run_test $old_ver "backward" - -echo " === Clear GreptimeDB data before running forward compatibility test" -rm -rf /tmp/greptimedb-standalone - -run_test $old_ver "forward" - -echo "Compatibility test run successfully!" diff --git a/tests/compat/util.sh b/tests/compat/util.sh deleted file mode 100755 index 94ec7ea3a2d..00000000000 --- a/tests/compat/util.sh +++ /dev/null @@ -1,82 +0,0 @@ -#!/bin/bash - -# Assemble the GreptimeDB binary download URL for a specific version. -binary_url() { - local ver="$1" - local bin_tar="greptime-$(uname -s | tr '[:upper:]' '[:lower:]')-amd64-v$ver.tar.gz" - echo "https://github.com/GreptimeTeam/greptimedb/releases/download/v$ver/$bin_tar" -} - -# Download a specific version of GreptimeDB binary tar file, untar it to folder `./bins/$ver`. -# `ver` is semver without prefix `v` -download_binary() { - local ver="$1" - local url="$(binary_url $ver)" - local bin_tar="greptime-$(uname -s | tr '[:upper:]' '[:lower:]')-amd64-v$ver.tar.gz" - - if [ -f ./bins/$ver/greptime ]; then - echo " === binaries exist: $(ls ./bins/$ver/* | tr '\n' ' ')" - chmod +x ./bins/$ver/* - return - fi - - if [ -f "$bin_tar" ]; then - echo " === tar file exists: $bin_tar" - else - echo " === Download binary ver: $ver" - echo " === Download binary url: $url" - curl --connect-timeout 5 --retry 5 --retry-delay 1 -L "$url" -o "$bin_tar" - fi - - mkdir -p ./bins/$ver - tar -xf "$bin_tar" --strip-components=1 -C ./bins/$ver - - echo " === unpacked: ./bins/$ver:" - ls -lh ./bins/$ver - - chmod +x ./bins/$ver/* -} - -# Test data compatibility that: -# - the data written by an old version of GreptimeDB can be read by the current one -# - the data written by the current version of GreptimeDB can be read by an old one ("forward" compatibility) -run_test() { - local old_ver="$1" - local forward="$2" - - local write_case_dir="./tests/compat/case/write" - local read_case_dir="./tests/compat/case/read" - local bin_old="./bins/$old_ver/greptime" - local bin_new="./bins/current/greptime" - local runner="./bins/current/sqlness-runner" - - echo " === Test with:" - echo " === old greptimedb version:" - "$bin_old" --version - echo " === new greptimedb version:" - "$bin_new" --version - - # "forward" means we are testing forward compatibility: - # the data generated by current version GreptimeDB can be used by old. - # So we run new GreptimeDB binary first to write, then run old to read. - # And the opposite for backward compatibility. - if [ "$forward" == "forward" ] - then - echo " === Running forward compat test ..." - echo " === Run test: write with current GreptimeDB" - $runner bare --bins-dir $(dirname $bin_new) --case-dir $write_case_dir - else - echo " === Running backward compat test ..." - echo " === Run test: write with old GreptimeDB" - $runner bare --bins-dir $(dirname $bin_old) --case-dir $write_case_dir - fi - - if [ "$forward" == 'forward' ] - then - echo " === Run test: read with old GreptimeDB" - $runner bare --bins-dir $(dirname $bin_old) --case-dir $read_case_dir - else - echo " === Run test: read with current GreptimeDB" - $runner bare --bins-dir $(dirname $bin_new) --case-dir $read_case_dir - fi -} diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index 542f10debfd..404b6332b94 100644 --- a/tests/compatibility/README.md +++ b/tests/compatibility/README.md @@ -2,7 +2,18 @@ Compatibility tests verify that one GreptimeDB version can restart on state written by another version. -Tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure. +SQL restart tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure. + +## Wire fixtures + +Frozen protobuf fixtures are exercised by the Rust tests that own their context +conversion, not by SQLness case discovery. The +[legacy query-context fixture](wire/read_preference/README.md) checks omitted +read-preference defaulting and unrelated metadata preservation: + +```shell +cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' +``` ## Quick Start diff --git a/tests/compatibility/cases/snapshot_parquet_copy/setup.sql b/tests/compatibility/cases/snapshot_parquet_copy/setup.sql index 8342e1847b5..38bc7c74aec 100644 --- a/tests/compatibility/cases/snapshot_parquet_copy/setup.sql +++ b/tests/compatibility/cases/snapshot_parquet_copy/setup.sql @@ -1,4 +1,6 @@ -CREATE TABLE snapshot_values (ts TIMESTAMP TIME INDEX, val BIGINT); -INSERT INTO snapshot_values VALUES (1, 42), (2, 7); +CREATE TABLE snapshot_values (ts TIMESTAMP TIME INDEX, val BIGINT, host STRING); +INSERT INTO snapshot_values VALUES (1, 42, ''), (2, 7, NULL); COPY snapshot_values TO '${SQLNESS_HOME}/snapshot_parquet_copy/values.parquet' WITH (FORMAT='parquet'); TRUNCATE TABLE snapshot_values; +CREATE TABLE snapshot_empty (ts TIMESTAMP TIME INDEX, val DOUBLE); +COPY snapshot_empty TO '${SQLNESS_HOME}/snapshot_parquet_copy/empty.parquet' WITH (FORMAT='parquet'); diff --git a/tests/compatibility/cases/snapshot_parquet_copy/verify.result b/tests/compatibility/cases/snapshot_parquet_copy/verify.result index d3e4d2483ab..6c590e1a1ad 100644 --- a/tests/compatibility/cases/snapshot_parquet_copy/verify.result +++ b/tests/compatibility/cases/snapshot_parquet_copy/verify.result @@ -4,9 +4,30 @@ Affected Rows: 2 SELECT * FROM snapshot_values ORDER BY ts; -+-------------------------+-----+ -| ts | val | -+-------------------------+-----+ -| 1970-01-01T00:00:00.001 | 42 | -| 1970-01-01T00:00:00.002 | 7 | -+-------------------------+-----+ ++-------------------------+-----+------+ +| ts | val | host | ++-------------------------+-----+------+ +| 1970-01-01T00:00:00.001 | 42 | | +| 1970-01-01T00:00:00.002 | 7 | | ++-------------------------+-----+------+ + +SELECT val, host IS NULL AS host_is_null, host = '' AS host_is_empty FROM snapshot_values ORDER BY ts; + ++-----+--------------+---------------+ +| val | host_is_null | host_is_empty | ++-----+--------------+---------------+ +| 42 | false | true | +| 7 | true | | ++-----+--------------+---------------+ + +COPY snapshot_empty FROM '${SQLNESS_HOME}/snapshot_parquet_copy/empty.parquet' WITH (FORMAT='parquet'); + +Affected Rows: 0 + +SELECT COUNT(*) FROM snapshot_empty; + ++----------+ +| count(*) | ++----------+ +| 0 | ++----------+ diff --git a/tests/compatibility/cases/snapshot_parquet_copy/verify.sql b/tests/compatibility/cases/snapshot_parquet_copy/verify.sql index c782dcacb1b..0e35e3c127f 100644 --- a/tests/compatibility/cases/snapshot_parquet_copy/verify.sql +++ b/tests/compatibility/cases/snapshot_parquet_copy/verify.sql @@ -1,2 +1,5 @@ COPY snapshot_values FROM '${SQLNESS_HOME}/snapshot_parquet_copy/values.parquet' WITH (FORMAT='parquet'); SELECT * FROM snapshot_values ORDER BY ts; +SELECT val, host IS NULL AS host_is_null, host = '' AS host_is_empty FROM snapshot_values ORDER BY ts; +COPY snapshot_empty FROM '${SQLNESS_HOME}/snapshot_parquet_copy/empty.parquet' WITH (FORMAT='parquet'); +SELECT COUNT(*) FROM snapshot_empty; diff --git a/tests/compatibility/wire/read_preference/README.md b/tests/compatibility/wire/read_preference/README.md new file mode 100644 index 00000000000..9e7b9ba74c1 --- /dev/null +++ b/tests/compatibility/wire/read_preference/README.md @@ -0,0 +1,45 @@ +# Legacy query-context wire compatibility + +`legacy_query_context.pb` is a generated legacy-shaped protobuf fixture, not a +capture from a released binary. It uses the unchanged `greptime.v1.QueryContext` +schema at greptime-proto revision +`9533078d6fd4bdaddbca0750cb9bf7ad64b2e568` and omits the internal +`query.read_preference` extension. + +Logical contents: + +- Catalog `c1`, schema `s1`, timezone `UTC`. +- Unrelated extension `flow.return_region_seq = true`. +- Channel `Internal` (`255`); no snapshot sequences or explain options. + +Run the real protobuf decode and session conversion regression: + +```shell +cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' +``` + +The test checks that the omitted preference defaults to Leader, that the original +fields and unrelated extension survive conversion and re-encoding, and that +Leader does not introduce a preference carrier. Existing session tests cover +present invalid or unsupported values. This fixture does not demonstrate that +an old receiver honors a non-Leader preference sent by an updated node. + +## Reproduce the fixture + +With `PROTO_ROOT` pointing to the `proto` directory of the revision above: + +```shell +printf '%s\n' \ + 'current_catalog: "c1"' \ + 'current_schema: "s1"' \ + 'timezone: "UTC"' \ + 'extensions { key: "flow.return_region_seq" value: "true" }' \ + 'channel: 255' | + protoc --proto_path="$PROTO_ROOT" --encode=greptime.v1.QueryContext \ + greptime/v1/common.proto > legacy_query_context.pb +``` + +Generated with `libprotoc 25.1`. The 48-byte fixture has SHA-256 +`4127ebbee8cdb1375ef8ad52754be56ed206c12513e5f61c4354c04587cba638`. +The test never rewrites it and does not require protobuf map ordering to remain +stable. diff --git a/tests/compatibility/wire/read_preference/legacy_query_context.pb b/tests/compatibility/wire/read_preference/legacy_query_context.pb new file mode 100644 index 00000000000..1990ede9411 --- /dev/null +++ b/tests/compatibility/wire/read_preference/legacy_query_context.pb @@ -0,0 +1,3 @@ + +c1s1"UTC* +flow.return_region_seqtrue0 \ No newline at end of file diff --git a/tests/perf/query_cases/promql_long_range_rate/case.toml b/tests/perf/query_cases/promql_long_range_rate/case.toml new file mode 100644 index 00000000000..51dfe1f1a86 --- /dev/null +++ b/tests/perf/query_cases/promql_long_range_rate/case.toml @@ -0,0 +1,119 @@ +# PromQL long-window rate benchmark for issue #9070. +# +# 2048 series × 1440 samples at 15-second cadence = 2,949,120 rows spanning the +# six hours from 2024-01-01T00:00:00Z to 2024-01-01T06:00:00Z. Evaluation covers +# the last two hours of that data (04:00:00Z–06:00:00Z at a 15-second step, 481 +# evaluations), so every evaluation holds a full 240-sample one-hour window. +# `rate_1h` is the primary regressed shape: it drives the RangeManipulate +# timestamp payload and the cursor scan across a long, heavily overlapping +# window set. `rate_5m` is the small-window control, `offset 1h` exercises the +# shifted payload path, and the bare selector controls unrelated-path variance. +# +# The generator derives each tag value as `series % cardinality`, so the number +# of distinct (host, instance) series is the least common multiple of the two +# cardinalities. `instance` declares 2048 values so the host(16) × instance tag +# shape still yields 2048 distinct series; a 128-value instance tag would cap +# this fixture at 128 distinct series. +# +# The deterministic wave is not a true monotonic counter, so the rate queries +# measure latency on the regular rate path rather than a numeric counter-rate +# oracle. + +[case] +name = "promql_long_range_rate" +description = "PromQL long-window rate latency regression for #9070 over six hours of 15-second samples" +issue = "https://github.com/GreptimeTeam/greptimedb/issues/9070" + +[scenario] +kind = "direct_readable_sst" +seed = 9070 + +[[scenario.tables]] +database = "public" +name = "promql_long_range_rate" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "instance"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "instance" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 2048, prefix = "instance" } + +[[scenario.tables.columns]] +name = "value" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 1000.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +# Eight 45-minute SSTs (180 timestamps each) tile the six-hour range. +# A (03:00:00Z, 06:00:00Z] scan prunes to four of them; the offset query scans +# the five SSTs covering [02:00:00Z, 05:00:00Z). +[scenario.layout] +regions = 1 +sst_count = 8 +rows_per_sst = 368640 +row_group_size = 36864 +series_count = 2048 +start_unix_nanos = 1704067200000000000 +step_nanos = 15000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +# 240 samples per full window: the primary regressed shape. +[[scenario.queries]] +name = "rate_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[1h])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# Twenty samples per full window: small-window control. +[[scenario.queries]] +name = "rate_5m" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[5m])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# The same long window shifted back an hour: non-zero-offset payload path. +[[scenario.queries]] +name = "rate_1h_offset_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[1h] offset 1h)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# Bare selector over the same evaluation range: unrelated-path variance control. +[[scenario.queries]] +name = "plain_selector_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') promql_long_range_rate{host=~'host.*'}" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 diff --git a/tests/perf/test_agent_observability_summary.py b/tests/perf/test_agent_observability_summary.py index 8223231c908..a6c1774ba0f 100644 --- a/tests/perf/test_agent_observability_summary.py +++ b/tests/perf/test_agent_observability_summary.py @@ -14,11 +14,13 @@ import importlib.util import json +import sys import tempfile import unittest from pathlib import Path SCRIPT = Path(__file__).parents[2] / '.github/scripts/agent-observability-summary.py' +sys.path.insert(0, str(SCRIPT.parent)) spec = importlib.util.spec_from_file_location('observability_summary', SCRIPT) summary = importlib.util.module_from_spec(spec) spec.loader.exec_module(summary) diff --git a/tests/perf/test_aliyun_ecs_runner_scripts.py b/tests/perf/test_aliyun_ecs_runner_scripts.py index 53257da8732..652eae77cfe 100644 --- a/tests/perf/test_aliyun_ecs_runner_scripts.py +++ b/tests/perf/test_aliyun_ecs_runner_scripts.py @@ -30,6 +30,9 @@ from types import SimpleNamespace from unittest.mock import Mock, patch SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" +sys.path.insert(0, str(SCRIPTS_DIR)) + +import runner_utils def load_module(name: str, filename: str): @@ -58,7 +61,7 @@ class GitHubApiResponseTest(unittest.TestCase): response.read.return_value = body response.__enter__ = Mock(return_value=response) response.__exit__ = Mock(return_value=False) - with patch.object(provision.urllib.request, "urlopen", return_value=response): + with patch.object(runner_utils.urllib.request, "urlopen", return_value=response): self.assertEqual(provision.github_api("token", "DELETE", "/test"), expected) def test_invalid_json_still_fails(self): @@ -66,10 +69,26 @@ class GitHubApiResponseTest(unittest.TestCase): response.read.return_value = b"not json" response.__enter__ = Mock(return_value=response) response.__exit__ = Mock(return_value=False) - with patch.object(provision.urllib.request, "urlopen", return_value=response): + with patch.object(runner_utils.urllib.request, "urlopen", return_value=response): with self.assertRaises(ValueError): provision.github_api("token", "GET", "/test") + def test_http_errors_still_exit_instead_of_becoming_recoverable(self): + error = runner_utils.urllib.error.HTTPError( + "https://api.github.com/test", 403, "Forbidden", {}, io.BytesIO(b"denied") + ) + with patch.object(runner_utils.urllib.request, "urlopen", side_effect=error): + with self.assertRaisesRegex(SystemExit, "HTTP 403: denied"): + provision.github_api("token", "GET", "/test") + + def test_unregister_preserves_legacy_failure_semantics(self): + with patch.object(runner_utils, "find_runner_by_name", return_value={"id": 42}): + with patch.object(runner_utils, "github_api", side_effect=SystemExit("HTTP 403")): + with self.assertRaisesRegex(SystemExit, "HTTP 403"): + runner_utils.deregister_runner("token", "owner/repo", "runner") + with patch.object(runner_utils, "github_api", side_effect=RuntimeError("connection lost")): + self.assertFalse(runner_utils.deregister_runner("token", "owner/repo", "runner")) + class ProvisionNamingTest(unittest.TestCase): def test_runner_name_and_label_derive_from_run_id(self) -> None: diff --git a/tests/perf/test_benchmark_version_comparison.py b/tests/perf/test_benchmark_version_comparison.py new file mode 100644 index 00000000000..56c2619e81c --- /dev/null +++ b/tests/perf/test_benchmark_version_comparison.py @@ -0,0 +1,138 @@ +# 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. + +import copy +import importlib.util +import json +from pathlib import Path +import tempfile +import unittest + +SCRIPTS = Path(__file__).resolve().parents[2] / '.github/scripts' +spec = importlib.util.spec_from_file_location('comparison', SCRIPTS / 'benchmark-version-comparison.py') +comparison = importlib.util.module_from_spec(spec) +spec.loader.exec_module(comparison) + + +def write(path, value): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value)) + + +class VersionComparison(unittest.TestCase): + def test_opt_in_validation(self): + comparison.validate({}) + comparison.validate({'COMPARE_GREPTIMEDB': 'false', 'GREPTIMEDB_TAG': 'invalid\n'}) + good = dict(COMPARE_GREPTIMEDB='true', RUN_GREPTIMEDB='true', + BASELINE_GREPTIMEDB_TAG='v1.2.1', GREPTIMEDB_TAG='v1.3.0-beta.1') + comparison.validate(good) + for key, value in [('RUN_GREPTIMEDB', 'false'), ('BASELINE_GREPTIMEDB_TAG', ''), + ('BASELINE_GREPTIMEDB_TAG', 'v1.3.0-beta.1'), + ('BASELINE_GREPTIMEDB_TAG', 'a,image=other'), ('GREPTIMEDB_TAG', 'x\ny'), + ('COMPARE_GREPTIMEDB', 'yes')]: + with self.subTest(key=key, value=value), self.assertRaises(ValueError): + comparison.validate(good | {key: value}) + + def test_improvement_direction(self): + self.assertAlmostEqual(comparison.improvement(100, 80), 20) + self.assertAlmostEqual(comparison.improvement(100, 120), -20) + self.assertAlmostEqual(comparison.improvement(100, 120, True), 20) + for old, new in [(0, 3), (None, 4), (1, -1), (1, float('nan')), (True, 2)]: + self.assertIsNone(comparison.improvement(old, new)) + + def test_agent_result_gates(self): + summary = dict(passed=True, contract_sha256='contract', data_profile='S', query_window=[1, 2], + warmup_runs=1, measured_runs=5, per_query=[dict(query_id='Q01', result_fingerprint='answer', + success=5, errors=0, timeouts=0, p50_ms=100, p95_ms=100, p99_ms=100)]) + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + base = root / 'baseline/greptimedb-1/greptimedb-1/measured-summary.json' + candidate = root / 'greptimedb-1/greptimedb-1/measured-summary.json' + for r in (root / 'baseline', root): + (r / 'greptimedb-1').mkdir(parents=True, exist_ok=True) + (r / 'greptimedb-1/exit-code.txt').write_text('0\n') + write(base, summary) + changed = copy.deepcopy(summary); changed['per_query'][0]['p50_ms'] = 80 + write(candidate, changed) + self.assertAlmostEqual(comparison.agent(root)[0]['improvement_percent'], 20) + for key, value in [('errors', 1), ('timeouts', 1), ('result_fingerprint', 'wrong'), ('success', 0)]: + broken = copy.deepcopy(changed); broken['per_query'][0][key] = value + write(candidate, broken) + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + (root / 'greptimedb-1/exit-code.txt').write_text('1\n') + write(candidate, changed) + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + (root / 'greptimedb-1/exit-code.txt').write_text('0\n') + changed['query_window'] = [2, 3]; write(candidate, changed) + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + candidate.unlink() + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + + def test_traces_uses_measured_only_and_rejects_failed_query(self): + answer = dict(query_id='T01', outcome='passed', fingerprint='answer', row_count=1, elapsed_ms=100) + sample = dict(answer, phase='measured') + record = dict(success=True, expected={'contracts': {'T01': {'fingerprint': 'answer', 'row_count': 1}}}, + protocol={'runs': 5}, audit_expected={}, corpus='sanity', + execution=dict(db_cpus='8', db_memory='16g', load_workers_requested=4, max_in_flight=8), + correctness=[answer], samples=[dict(sample, phase='warmup', elapsed_ms=99999), sample, sample]) + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp);base = root / 'baseline/runs/greptimedb/run.json'; candidate = root / 'runs/greptimedb/run.json' + write(base, record); changed = copy.deepcopy(record) + for s in changed['samples']: + if s['phase'] == 'measured': s['elapsed_ms'] = 80 + write(candidate, changed) + self.assertAlmostEqual(comparison.traces(root)[0]['improvement_percent'], 20) + self.assertIsNone(comparison.traces(root)[1]['improvement_percent']) + for key, value in [('outcome', 'failed'), ('fingerprint', 'wrong'), ('row_count', 2)]: + broken = copy.deepcopy(changed); broken['samples'][1][key] = value; write(candidate, broken) + self.assertIsNone(comparison.traces(root)[0]['improvement_percent']) + + def test_long_range_gates(self): + summary = dict(workload='vmbench-long-range', query_file_sha256='hash', start='start', end='end', step='30m', + profile='small', data_duration='32d', scrape_interval='30s', cache_policy_effective='warm', + mode='serial', connection_mode='force-close', request_timeout_seconds=120, tries=5, warmup_runs=1, + errors=0, timeouts=0, per_query=[dict(query_index=1, query='up', series=1, points=2, + success=5, errors=0, timeouts=0, p50_ms=100)]) + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp); base = root/'greptimedb-baseline/results-vmbench-long-range/serial-summary.json' + candidate = root/'greptimedb/results-vmbench-long-range/serial-summary.json' + health = {phase: {target: {'state': 'Running=true ExitCode=0 OOMKilled=false'} + for target in ('greptimedb-baseline', 'greptimedb')} + for phase in ('post_load_results', 'post_query_results')} + health['post_load_results'] = {k: {'docker': v} for k, v in health['post_load_results'].items()} + write(root / 'run-manifest.json', health) + write(base, summary); write(candidate, summary) + self.assertEqual(comparison.long_range(root)[0]['improvement_percent'], 0) + for key, value in [('start','different'),('tries',1),('query_file_sha256','other'),('errors',1)]: + write(candidate, summary | {key:value}) + self.assertIsNone(comparison.long_range(root)[0]['improvement_percent']) + write(candidate, summary) + broken=copy.deepcopy(health);broken['post_query_results']['greptimedb']['state']='Running=false ExitCode=137 OOMKilled=true' + write(root / 'run-manifest.json', broken) + self.assertIsNone(comparison.long_range(root)[0]['improvement_percent']) + write(root / 'run-manifest.json', health) + empty=copy.deepcopy(summary);empty['per_query'][0]['points']=0;write(candidate,empty) + self.assertIsNone(comparison.long_range(root)[0]['improvement_percent']) + + def test_missing_results_and_rendering(self): + with tempfile.TemporaryDirectory() as tmp: + for load in (comparison.agent, comparison.traces, comparison.long_range): + self.assertTrue(all(r['improvement_percent'] is None for r in load(Path(tmp)))) + rendered = comparison.render(dict(baseline_tag='old', candidate_tag='', rows=[])) + self.assertIn('<new|tag>', rendered) + self.assertIn('Results unavailable', rendered) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/perf/test_ci_runner_scripts.py b/tests/perf/test_ci_runner_scripts.py new file mode 100644 index 00000000000..c703e41c62b --- /dev/null +++ b/tests/perf/test_ci_runner_scripts.py @@ -0,0 +1,616 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline regression coverage for CI runner selection and lifecycle.""" + +import importlib.util +import json +import os +import re +import subprocess +import sys +import tempfile +import textwrap +import unittest +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +SCRIPTS = Path(__file__).parents[2] / ".github/scripts" +sys.path.insert(0, str(SCRIPTS)) + + +def load(name, filename): + spec = importlib.util.spec_from_file_location(name, SCRIPTS / filename) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +config = load("ci_config_test", "ci-runner-config.py") +aws = load("aws_provision_test", "aws-ec2-runner-provision.py") +teardown = load("aws_teardown_test", "aws-ec2-runner-teardown.py") + + +def provision_config(**overrides): + return aws.ProvisionConfig(**({ + "region": "ap-southeast-1", "image_id": "ami-test", "instance_type": "c7i.2xlarge", + "subnet_id": "subnet-test", "security_group_id": "sg-test", "repo": "owner/repo", + "run_id": "123-2", "github_token": "secret", "system_disk_gib": 80, + } | overrides)) + + +def client_with_image(size=40, arch="x86_64"): + client = Mock() + client.describe_images.return_value = {"Images": [{ + "State": "available", "Architecture": arch, "RootDeviceName": "/dev/sda1", + "BlockDeviceMappings": [{"DeviceName": "/dev/sda1", "Ebs": {"VolumeSize": size}}], + }]} + client.describe_instance_types.return_value = {"InstanceTypes": [{ + "ProcessorInfo": {"SupportedArchitectures": ["x86_64"]}, + }]} + client.run_instances.return_value = {"Instances": [{"InstanceId": "i-test"}]} + client.get_console_output.return_value = {"Output": "boot log"} + return client + + +class ProviderConfigTest(unittest.TestCase): + def test_defaults_and_overrides(self): + for provider, requested, expected in ( + ("Aliyun", "auto", "ecs.c9i.2xlarge"), ("AWS", "auto", "c7i.2xlarge"), + ("Aliyun", "ecs.c9i.4xlarge", "ecs.c9i.4xlarge"), + ("AWS", "c6i.4xlarge", "c6i.4xlarge"), + ("AWS", "c7i.metal-24xl", "c7i.metal-24xl"), + ): + with self.subTest(provider=provider, requested=requested): + self.assertEqual(config.resolve_instance_type(provider, requested), expected) + + def test_traces_workflow_runner_defaults_and_overrides(self): + workflow = (SCRIPTS.parent / "workflows/tracesbench.yml").read_text() + step = workflow.split(" - name: Resolve CI runner\n", 1)[1].split(" - name:", 1)[0] + command = textwrap.dedent(step.split(" run: |\n", 1)[1]) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "greptimedb").symlink_to(SCRIPTS.parents[1], target_is_directory=True) + output = root / "output" + for provider, requested, expected in ( + ("Aliyun", "auto", "ecs.g9i.2xlarge"), + ("AWS", "auto", "m7i.2xlarge"), + ("Aliyun", "ecs.c9i.4xlarge", "ecs.c9i.4xlarge"), + ("AWS", "c7i.4xlarge", "c7i.4xlarge"), + ("invalid", "auto", None), + ("AWS", "ecs.g9i.2xlarge", None), + ): + with self.subTest(provider=provider, requested=requested): + output.write_text("") + result = subprocess.run( + ["bash", "-c", command], cwd=root, capture_output=True, text=True, + env=dict(os.environ, PROVIDER=provider, REQUESTED_INSTANCE_TYPE=requested, + GITHUB_OUTPUT=str(output)), timeout=10, + ) + if expected is None: + self.assertNotEqual(result.returncode, 0) + self.assertEqual(output.read_text(), "") + else: + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(output.read_text(), f"instance_type={expected}\n") + self.assertRegex(workflow, r"(?s)db_cpus:.*?default: '8'") + self.assertRegex(workflow, r"(?s)db_memory:.*?default: 32g") + + def test_traces_selected_targets_and_image_tag(self): + workflow = (SCRIPTS.parent / "workflows/tracesbench.yml").read_text() + step = workflow.split(" - name: Validate selected targets\n", 1)[1].split(" - name:", 1)[0] + command = textwrap.dedent(step.split(" run: |\n", 1)[1]) + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) / "output" + for enabled, tag, expected in ( + (("true", "false", "false"), "latest", "greptimedb"), + (("false", "true", "false"), "", "victoriatraces"), + (("false", "false", "true"), "", "tempo"), + (("true", "true", "true"), "v1.3.0-beta.1", "greptimedb,victoriatraces,tempo"), + (("false", "false", "false"), "latest", None), + (("true", "false", "false"), "latest; touch /tmp/unexpected", None), + (("true", "false", "false"), "", None), + (("invalid", "false", "false"), "latest", None), + ): + with self.subTest(enabled=enabled, tag=tag): + output.write_text("") + env = dict(os.environ, GREPTIMEDB_TAG=tag, GITHUB_OUTPUT=str(output)) + env.update(zip(("greptimedb", "victoriatraces", "tempo"), enabled)) + result = subprocess.run(["bash", "-c", command], env=env, + capture_output=True, text=True, timeout=5) + if expected is None: + self.assertNotEqual(result.returncode, 0) + self.assertNotIn("targets=", output.read_text()) + else: + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn(f"targets={expected}\n", output.read_text()) + + def test_wrong_provider_and_shell_fragments_rejected(self): + for provider, requested in (("ec2", "auto"), ("AWS", "ecs.c9i.2xlarge"), + ("Aliyun", "c7i.2xlarge"), ("AWS", ""), + ("AWS", "c7i.2xlarge; touch /tmp/unexpected")): + with self.subTest(provider=provider, requested=requested), self.assertRaises(ValueError): + config.resolve_instance_type(provider, requested) + + def test_cli_writes_resolved_output(self): + with tempfile.TemporaryDirectory() as tmp: + output = Path(tmp) / "output" + result = subprocess.run([sys.executable, str(SCRIPTS / "ci-runner-config.py"), + "--provider", "AWS"], + env=os.environ | {"GITHUB_OUTPUT": str(output)}, + capture_output=True, text=True) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(output.read_text(), "instance_type=c7i.2xlarge\n") + + + +class BenchmarkSetupTest(unittest.TestCase): + def run_setup(self, missing_tools="", docker_installed=True, docker_ready=True): + action = (SCRIPTS.parent / "actions/setup-benchmark/action.yml").read_text() + steps = re.findall(r" run: \|\n((?: .*\n|\n)+)", action) + self.assertEqual(len(steps), 2) + script = "\n".join(textwrap.dedent(step) for step in steps) + # Execute the actual action twice with host commands stubbed: installation + # changes availability, so the second pass must perform no mutations. + prelude = r""" +function . { ID=ubuntu; VERSION_CODENAME=noble; } +uname() { echo x86_64; } +dpkg() { echo amd64; } +curl() { echo fixture-key; } +command() { + if [[ "$1" == -v ]]; then + if [[ "$2" == docker ]]; then [[ "$DOCKER_INSTALLED" == true ]]; return; fi + if [[ " $MISSING_TOOLS " == *" $2 "* ]]; then return 1; fi + return 0 + fi + builtin command "$@" +} +docker() { [[ "$DOCKER_READY" == true ]]; } +sudo() { + printf '%s\n' "$*" >> "$CALLS" + if [[ "$*" == *' apt-get install '* ]]; then + MISSING_TOOLS='' + [[ "$*" != *' docker-ce '* ]] || DOCKER_INSTALLED=true + fi + if [[ "$*" == *'systemctl enable --now docker' ]]; then DOCKER_READY=true; fi + if [[ "$*" == *' tee '* ]]; then while IFS= read -r line; do :; done; fi + return 0 +} +""" + with tempfile.TemporaryDirectory() as tmp: + calls = Path(tmp) / "calls" + calls.touch() + env = os.environ | {"CALLS": str(calls), "MISSING_TOOLS": missing_tools, + "DOCKER_INSTALLED": str(docker_installed).lower(), + "DOCKER_READY": str(docker_ready).lower()} + result = subprocess.run(["bash", "-c", prelude + script + + '\nprintf "SECOND PASS\\n" >> "$CALLS"\n' + script], + env=env, capture_output=True, text=True) + self.assertEqual(result.returncode, 0, result.stderr) + first, second = calls.read_text().split("SECOND PASS\n") + self.assertEqual(second, "", "Second setup must reuse installed dependencies") + return first + + def test_prepared_runner_skips_all_privileged_changes(self): + self.assertEqual(self.run_setup(), "") + + def test_missing_tools_are_installed_without_changing_docker(self): + calls = self.run_setup(missing_tools="jq rsync") + self.assertIn("--no-install-recommends jq rsync", calls) + self.assertIn("NEEDRESTART_MODE=l", calls) + self.assertNotIn("docker", calls) + + def test_missing_docker_is_installed_once(self): + calls = self.run_setup(docker_installed=False, docker_ready=False) + self.assertIn("NEEDRESTART_MODE=l", calls) + self.assertIn("docker-ce docker-ce-cli containerd.io docker-buildx-plugin", calls) + self.assertIn("systemctl enable --now docker", calls) + + def test_existing_stopped_docker_is_started_without_installation(self): + self.assertEqual(self.run_setup(docker_ready=False), "-n systemctl enable --now docker\n") + + +class AwsProvisionTest(unittest.TestCase): + def test_host_policy_matches_aliyun_before_runner_activation(self): + aliyun = load("aliyun_host_policy_test", "aliyun-ecs-runner-provision.py") + ecs_script = aliyun.render_user_data("name", "label", "token", "owner/repo") + aws_script = aws.render_user_data("name", "token", "owner/repo", "https://example.com/runner.tgz") + start = ecs_script.index("# Mask systemd-oomd") + end = ecs_script.index("free --human", start) + len("free --human") + policy = ecs_script[start:end] + self.assertIn(policy, aws_script) + self.assertLess(aws_script.index(policy), aws_script.index("./svc.sh start")) + self.assertLess(aws_script.index("OOMPolicy=continue"), aws_script.index("./svc.sh start")) + self.assertLess(aws_script.index("./svc.sh install runner"), aws_script.index("service_name=$(cat .service)")) + self.assertIn('"/etc/systemd/system/$service_name.d/oom.conf"', aws_script) + + def test_ubuntu_image_resolution_and_override(self): + client = Mock() + client.describe_images.return_value = {"Images": [ + {"ImageId": "ami-new", "CreationDate": "2026-09-04", "Name": "noble-new"}, + {"ImageId": "ami-old", "CreationDate": "2026-08-04", "Name": "noble-old"}, + ]} + self.assertEqual(aws.resolve_image(client, "auto"), "ami-new") + request = client.describe_images.call_args.kwargs + self.assertEqual(request["Owners"], ["099720109477"]) + self.assertIn({"Name": "name", "Values": [ + "ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*"]}, request["Filters"]) + client.reset_mock() + self.assertEqual(aws.resolve_image(client, "ami-pinned"), "ami-pinned") + client.describe_images.assert_not_called() + client.describe_images.return_value = {"Images": []} + with self.assertRaisesRegex(ValueError, "No Canonical Ubuntu"): + aws.resolve_image(client, "auto") + + def test_disk_size_and_architecture_fail_before_allocation(self): + for client, pattern in ((client_with_image(size=500), "smaller than AMI"), + (client_with_image(arch="arm64"), "x86_64")): + with self.subTest(pattern=pattern), patch.object(aws, "create_registration_token") as token: + with self.assertRaisesRegex(ValueError, pattern): + aws.run_instance(client, provision_config()) + client.run_instances.assert_not_called() + token.assert_not_called() + + def test_arm_instance_rejected_before_registration(self): + client = client_with_image() + client.describe_instance_types.return_value["InstanceTypes"][0]["ProcessorInfo"]["SupportedArchitectures"] = ["arm64"] + with patch.object(aws, "create_registration_token") as token: + with self.assertRaisesRegex(ValueError, "x86_64"): + aws.run_instance(client, provision_config(instance_type="r8g.2xlarge")) + token.assert_not_called() + client.run_instances.assert_not_called() + + def test_create_request_and_recovery_tags(self): + client = client_with_image() + with patch.object(aws, "create_registration_token", return_value="short-lived-token"), \ + patch.object(aws, "runner_download", return_value="https://example.com/runner.tar.gz"): + self.assertEqual(aws.run_instance(client, provision_config(system_disk_gib=100)), "i-test") + request = client.run_instances.call_args.kwargs + self.assertEqual(request["BlockDeviceMappings"], [{"DeviceName": "/dev/sda1", "Ebs": { + "VolumeSize": 100, "VolumeType": "gp3", "Iops": 3000, "Throughput": 125, + "DeleteOnTermination": True}}]) + self.assertEqual(request["ClientToken"], "ci-ec2-123-2") + self.assertNotIn("InstanceMarketOptions", request) + self.assertNotIn("IamInstanceProfile", request) + for specification in request["TagSpecifications"]: + tags = {t["Key"]: t["Value"] for t in specification["Tags"]} + self.assertEqual(tags["managed-by"], "query-regression-ci") + self.assertEqual(tags[aws.REPO_TAG], "owner/repo") + self.assertEqual(tags["query-regression-run-id"], "123-2") + self.assertEqual(tags["runner-ttl-hours"], "8") + self.assertNotIn("secret", request["UserData"]) + self.assertIn("--token short-lived-token", request["UserData"]) + self.assertNotIn("apt-get install", request["UserData"]) + self.assertIn('for tool in git curl tar sudo; do command -v "$tool"; done', request["UserData"]) + self.assertLess(request["UserData"].index("usermod -aG docker runner"), + request["UserData"].index("./svc.sh start")) + self.assertIn("visudo -cf /etc/sudoers.d/ci-runner", request["UserData"]) + self.assertLess(request["UserData"].index("$nrconf{override_rc}"), + request["UserData"].index("./svc.sh start")) + check = subprocess.run(["bash", "-n"], input=request["UserData"], text=True, capture_output=True) + self.assertEqual(check.returncode, 0, check.stderr) + + def test_partial_provision_keeps_outputs_and_console(self): + client = client_with_image() + client.get_waiter.return_value.wait.side_effect = TimeoutError("not running") + with tempfile.TemporaryDirectory() as tmp: + output = Path(tmp) / "output" + with patch.dict(os.environ, {"GITHUB_OUTPUT": str(output)}), patch.object(aws, "run_instance", return_value="i-test"): + with self.assertRaises(TimeoutError): + aws.provision(client, provision_config()) + self.assertEqual(output.read_text().splitlines(), [ + "instance_id=i-test", "label=ci-ec2-123-2", + "runner_name=ci-ec2-123-2"]) + client.get_console_output.assert_called_once() + + def test_online_poll_and_timeout(self): + for online in (True, False): + with self.subTest(online=online), patch.object(aws, "run_instance", return_value="i-test"), \ + patch.object(aws, "find_runner_by_name", side_effect=[None, {"status": "online"}]), \ + patch.object(aws.time, "sleep"), \ + patch.object(aws.time, "monotonic", side_effect=[0, 1, 2 if online else 601]): + client = client_with_image() + if online: + self.assertEqual(aws.provision(client, provision_config()), 0) + else: + with self.assertRaises(TimeoutError): + aws.provision(client, provision_config()) + client.get_console_output.assert_called_once() + + +class AwsTeardownTest(unittest.TestCase): + NOW = datetime(2026, 9, 23, 10, tzinfo=timezone.utc) + + def setUp(self): + sleeper = patch.object(teardown.time, "sleep") + sleeper.start() + self.addCleanup(sleeper.stop) + + def instance(self, identity="i-test", age=9, ttl="8"): + tags = [{"Key": teardown.RUNNER_TAG, "Value": "ci-ec2-123-2"}] + if ttl is not None: + tags.append({"Key": teardown.TTL_TAG_KEY, "Value": ttl}) + return {"InstanceId": identity, "LaunchTime": self.NOW - timedelta(hours=age), "Tags": tags} + + def test_pagination_and_ownership_filters(self): + client = Mock() + client.get_paginator.return_value.paginate.return_value = [ + {"Reservations": [{"Instances": [self.instance("i-a")]}]}, + {"Reservations": [{"Instances": [self.instance("i-b")]}]}, + ] + self.assertEqual([i["InstanceId"] for i in teardown.managed_instances(client, "owner/repo", "123-2")], ["i-a", "i-b"]) + self.assertEqual(client.get_paginator.return_value.paginate.call_args.kwargs["Filters"], [ + {"Name": "tag:managed-by", "Values": ["query-regression-ci"]}, + {"Name": "tag:github-repository", "Values": ["owner/repo"]}, + {"Name": "tag:query-regression-run-id", "Values": ["123-2"]}, + ]) + + def test_sweep_expiry_and_missing_or_invalid_ttl(self): + instances = [self.instance("expired", age=8), self.instance("fresh", age=7), + self.instance("invalid", ttl="bad"), + self.instance("missing-expired", age=4, ttl=None), + self.instance("missing-fresh", age=3, ttl=None)] + with patch.object(teardown, "managed_instances", return_value=instances), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", now=self.NOW), 0) + self.assertEqual([c.args[1] for c in delete.call_args_list], ["expired", "missing-expired"]) + self.assertEqual(unregister.call_count, 2) + + def test_ttl_decisions_match_aliyun(self): + aliyun = load("aliyun_ttl_policy_test", "aliyun-ecs-runner-teardown.py") + for age, ttl in [(3, None), (4, None), (7, "8"), (8, "8"), (9, "8"), + (10, "bad"), (10, ""), (10, "0"), (200, "169"), (168, "168")]: + instance = self.instance(age=age, ttl=ttl) + created = instance["LaunchTime"].strftime("%Y-%m-%dT%H:%M:%SZ") + expected = aliyun.expired_instance_names( + [(instance["InstanceId"], "runner", created, ttl)], self.NOW, timedelta(hours=4)) + with self.subTest(age=age, ttl=ttl), \ + patch.object(teardown, "managed_instances", return_value=[instance]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True): + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", now=self.NOW), 0) + self.assertEqual([call.args[1] for call in delete.call_args_list], + [identity for identity, _ in expected]) + + def test_targeted_cleanup_recovers_missing_outputs_without_waiting_for_ttl(self): + with patch.object(teardown, "managed_instances", return_value=[self.instance(age=0)]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 0) + delete.assert_called_once() + unregister.assert_called_once_with("token", "owner/repo", "ci-ec2-123-2") + + def test_independent_cloud_and_github_cleanup_failures(self): + for cloud_ok, github_result in ((False, True), (True, RuntimeError("GitHub unavailable"))): + with self.subTest(cloud_ok=cloud_ok), \ + patch.object(teardown, "managed_instances", return_value=[self.instance()]), \ + patch.object(teardown, "delete_instance", return_value=cloud_ok) as delete, \ + patch.object(teardown, "deregister_runner", side_effect=[github_result]) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 1) + delete.assert_called_once() + unregister.assert_called_once() + + def test_already_gone_instance_still_unregisters(self): + with patch.object(teardown, "managed_instances", return_value=[]), \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 0) + unregister.assert_called_once() + + def test_retried_teardown_uses_original_instance_and_runner(self): + client = Mock() + client.get_paginator.return_value.paginate.return_value = [ + {"Reservations": [{"Instances": [self.instance()]}]}, + ] + with patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(client, "owner/repo", "token", run_id="123-3", + instance_id="i-test", runner_name="ci-ec2-123-2"), 0) + filters = client.get_paginator.return_value.paginate.call_args.kwargs["Filters"] + self.assertIn({"Name": "instance-id", "Values": ["i-test"]}, filters) + self.assertFalse(any(f["Name"] == "tag:query-regression-run-id" for f in filters)) + delete.assert_called_once_with(client, "i-test") + unregister.assert_called_once_with("token", "owner/repo", "ci-ec2-123-2") + client.get_paginator.return_value.paginate.return_value = [] + with patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(client, "owner/repo", "token", run_id="123-3", + instance_id="i-test", runner_name="ci-ec2-123-2"), 0) + unregister.assert_called_once_with("token", "owner/repo", "ci-ec2-123-2") + + def test_terminate_waits_and_handles_not_found(self): + client = client_with_image() + self.assertTrue(teardown.delete_instance(client, "i-test")) + client.get_waiter.assert_called_with("instance_terminated") + error = RuntimeError("gone") + error.response = {"Error": {"Code": "InvalidInstanceID.NotFound"}} + client.terminate_instances.side_effect = error + with patch.object(teardown.time, "sleep"): + self.assertTrue(teardown.delete_instance(client, "i-test")) + + def test_transient_not_found_is_retried(self): + client = client_with_image() + error = RuntimeError("not visible yet") + error.response = {"Error": {"Code": "InvalidInstanceID.NotFound"}} + client.terminate_instances.side_effect = [error, {}] + self.assertTrue(teardown.delete_instance(client, "i-test")) + self.assertEqual(client.terminate_instances.call_count, 2) + + def test_missing_tags_are_polled_before_cleanup(self): + with patch.object(teardown, "managed_instances", side_effect=[[], [self.instance()]]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True): + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 0) + delete.assert_called_once() + + def test_github_system_exit_does_not_skip_other_instances(self): + with patch.object(teardown, "managed_instances", return_value=[self.instance("i-a"), self.instance("i-b")]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", side_effect=[SystemExit("HTTP 403"), True]): + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 1) + self.assertEqual([c.args[1] for c in delete.call_args_list], ["i-a", "i-b"]) + +class BenchmarkImagePullTest(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.root = Path(self.tmp.name) + self.env = dict(os.environ, PATH=f"{self.root}:{os.environ['PATH']}", + PULL_TEST_ROOT=str(self.root)) + self.tool("docker", ''' + echo "$*" >> "$PULL_TEST_ROOT/docker.log" + count=$(wc -l < "$PULL_TEST_ROOT/docker.log") + codes=${PULL_TEST_CODES:-0} + code=$(printf '%s' "$codes" | cut -d, -f"$count") + exit "${code:-0}" + ''') + self.tool("timeout", ''' + echo "$*" >> "$PULL_TEST_ROOT/timeout.log" + shift 2 + if [[ ${PULL_TEST_TIMEOUT:-0} == 1 ]]; then exit 124; fi + exec "$@" + ''') + self.tool("sleep", 'echo "$*" >> "$PULL_TEST_ROOT/sleep.log"') + + def tool(self, name, body): + p = self.root / name + p.write_text("#!/bin/bash\nset -eu\n" + textwrap.dedent(body)) + p.chmod(0o755) + + def lines(self, name): + p = self.root / (name + ".log") + return p.read_text().splitlines() if p.exists() else [] + + def pull(self, codes, *args): + self.env["PULL_TEST_CODES"] = codes + return subprocess.run(["bash", str(SCRIPTS / "pull-benchmark-image.sh"), *args], + env=self.env, capture_output=True, text=True, timeout=5) + + def test_success_does_not_retry(self): + result = self.pull("0", "registry/test@sha256:" + "a" * 64) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(len(self.lines("docker")), 1) + self.assertEqual(self.lines("sleep"), []) + self.assertTrue(self.lines("timeout")[0].startswith("--kill-after=10s 180s docker pull ")) + + def test_transient_failure_then_success(self): + result = self.pull("1,1,0", "registry/test:tag") + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(self.lines("docker"), ["pull registry/test:tag"] * 3) + self.assertEqual(self.lines("sleep"), ["10", "20"]) + self.assertIn("attempt 3/3", result.stderr) + + def test_exhaustion_preserves_failure(self): + result = self.pull("1,1,7", "registry/test:tag") + self.assertEqual(result.returncode, 7) + self.assertEqual(len(self.lines("docker")), 3) + self.assertEqual(self.lines("sleep"), ["10", "20"]) + + def test_timeout_is_bounded_and_fatal(self): + self.env["PULL_TEST_TIMEOUT"] = "1" + result = self.pull("0", "registry/test:tag") + self.assertEqual(result.returncode, 124) + self.assertEqual(len(self.lines("timeout")), 3) + self.assertEqual(self.lines("sleep"), ["10", "20"]) + + def test_bad_arguments_do_not_call_docker(self): + for args in ([], [""], ["--help"], ["a", "b"]): + with self.subTest(args=args): + self.assertEqual(self.pull("0", *args).returncode, 2) + self.assertEqual(self.lines("docker"), []) + + def test_traces_prefetch_executes_catalog_refs_and_stops_on_failure(self): + traces = (SCRIPTS.parent / 'workflows/tracesbench.yml').read_text() + start = traces.index(' - name: Pull database images before generating data') + end = traces.index(' - name: Generate selected dataset once') + command = textwrap.dedent(traces[start:end].split(' run: |\n', 1)[1]) + helper = self.root / 'greptimedb/.github/scripts/pull-benchmark-image.sh' + helper.parent.mkdir(parents=True) + helper.write_text((SCRIPTS / helper.name).read_text()) + catalogs = self.root / 'o11ybench/workloads/tracesbench/catalogs' + catalogs.mkdir(parents=True) + images = [f'registry/{name}@sha256:' + 'a' * 64 + for name in ('greptimedb', 'victoriatraces', 'tempo')] + for name, image in zip(('greptimedb', 'victoriatraces', 'tempo'), images): + (catalogs / f'{name}.json').write_text(json.dumps({'target_identity': {'image': image}})) + self.tool('docker', ''' + echo "$*" >> "$PULL_TEST_ROOT/docker.log" + if [[ $1 == pull && $2 == ${PULL_FAIL_IMAGE:-} ]]; then exit 1; fi + if [[ $1 == image ]]; then printf '[{"RepoDigests":["greptime/greptimedb@sha256:%s"]}]\\n' "$(printf 'a%.0s' {1..64})"; fi + ''') + images[0] = 'greptime/greptimedb:v1.3.0-beta.1' + self.env.update(METADATA_ROOT=str(self.root), GREPTIMEDB_TAG='v1.3.0-beta.1', + GITHUB_ENV=str(self.root / 'env'), BASELINE_GREPTIMEDB_TAG='v1.2.1') + for selected, compare in ((('greptimedb',), False), (('tempo',), False), + (('greptimedb', 'victoriatraces', 'tempo'), False), + (('greptimedb',), True), (('greptimedb', 'victoriatraces', 'tempo'), True)): + with self.subTest(selected=selected): + self.env['TARGETS'] = ','.join(selected) + self.env['COMPARE_GREPTIMEDB'] = str(compare).lower() + (self.root / 'env').write_text('') + (self.root / 'docker.log').unlink(missing_ok=True) + result = subprocess.run(['bash', '-c', command], cwd=self.root, env=self.env, + capture_output=True, text=True, timeout=5) + self.assertEqual(result.returncode, 0, result.stderr) + selected_images = [image for target, image in zip(('greptimedb', 'victoriatraces', 'tempo'), images) + if target in selected] + if compare: + selected_images.append('greptime/greptimedb:v1.2.1') + self.assertEqual(self.lines('docker'), + [call for image in selected_images for call in (f'pull {image}', f'image inspect {image}')]) + expected_env = 'GREPTIMEDB_IMAGE=greptime/greptimedb@sha256:' + 'a' * 64 + '\n' + expected_env = expected_env if 'greptimedb' in selected else '' + if compare: + expected_env += 'BASELINE_GREPTIMEDB_IMAGE=greptime/greptimedb@sha256:' + 'a' * 64 + '\n' + self.assertEqual((self.root / 'env').read_text(), expected_env) + (self.root / 'docker.log').unlink() + self.env['PULL_FAIL_IMAGE'] = images[1] + result = subprocess.run(['bash', '-c', command], cwd=self.root, env=self.env, + capture_output=True, text=True, timeout=5) + self.assertNotEqual(result.returncode, 0) + self.assertEqual(self.lines('docker'), [f'pull {images[0]}', f'image inspect {images[0]}'] + + [f'pull {images[1]}'] * 3) + self.assertNotIn(f'pull {images[2]}', self.lines('docker')) + + def test_workflows_pull_before_generation(self): + workflows = SCRIPTS.parent / "workflows" + agent = (workflows / "agent-observability.yml").read_text() + traces = (workflows / "tracesbench.yml").read_text() + long_range = (workflows / "vmbench-long-range.yml").read_text() + for content in (agent, traces, long_range): + self.assertNotIn("docker pull ", content) + for reference in ("RUNTIME_IMAGE", "RESOLVED_RUNTIME_IMAGE", "reference"): + self.assertIn(f'pull-benchmark-image.sh "${reference}"', long_range) + self.assertLess(long_range.index('pull-benchmark-image.sh "$reference"'), + long_range.index('- name: Warm serial and concurrent benchmarks')) + self.assertIn('pull-benchmark-image.sh "$reference"', agent) + self.assertLess(agent.index('resolve_image runtime'), agent.index('- name: Generate dataset once')) + self.assertIn('pull-benchmark-image.sh "$RUNTIME_IMAGE"', traces) + self.assertIn('pull-benchmark-image.sh "$RESOLVED_RUNTIME_IMAGE"', traces) + start = traces.index(' - name: Pull database images before generating data') + end = traces.index(' - name: Generate selected dataset once') + self.assertLess(start, end) + step = traces[start:end] + self.assertIn('IFS=, read -ra targets <<< "$TARGETS"', step) + self.assertIn('for target in "${targets[@]}"', step) + self.assertIn('o11ybench/workloads/tracesbench/catalogs/$target.json', step) + self.assertIn('pull-benchmark-image.sh "$image"', step) + self.assertIn('set -euo pipefail', step) + self.assertNotIn('continue-on-error', step) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/perf/test_tracesbench_summary.py b/tests/perf/test_tracesbench_summary.py new file mode 100644 index 00000000000..cb7c4d8e06e --- /dev/null +++ b/tests/perf/test_tracesbench_summary.py @@ -0,0 +1,54 @@ +# 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. + +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path + +SCRIPT = Path(__file__).parents[2] / '.github/scripts/tracesbench-summary.py' +sys.path.insert(0, str(SCRIPT.parent)) +spec = importlib.util.spec_from_file_location('tracesbench_summary', SCRIPT) +summary = importlib.util.module_from_spec(spec) +spec.loader.exec_module(summary) + + +class LifecycleTest(unittest.TestCase): + def test_ingestion_layouts(self): + metrics = {'wall_seconds': 0.121723417, 'successful_requests': 17, 'failed_requests': 0} + cases = [ + (target, {'ingest': metrics}, ['121.7 ms', '17', '0']) + for target in summary.TARGETS + ] + [ + ('tempo', {'transport': {'ingest': metrics}}, ['121.7 ms', '17', '0']), + ('tempo', {'ingest': {'wall_seconds': 0, 'successful_requests': 0, 'failed_requests': 0}, + 'transport': {'ingest': metrics}}, ['0 µs', '0', '0']), + ('tempo', {'ingest': {}, 'transport': {'ingest': metrics}}, ['—', '—', '—']), + ('tempo', {'transport': {}}, ['—', '—', '—']), + ] + [(target, {}, ['—', '—', '—']) for target in summary.TARGETS] + [ + (target, {'transport': {'ingest': metrics}}, ['—', '—', '—']) + for target in ('greptimedb', 'victoriatraces') + ] + with tempfile.TemporaryDirectory() as tmp: + for target, load, expected in cases: + with self.subTest(target=target, load=load): + text = summary.lifecycle({}, [{'target': target, 'load': load}], Path(tmp), [target]) + row = next(line for line in text.splitlines() if line.startswith('| ' + target + ' |')) + cells = [cell.strip() for cell in row.split('|')[1:-1]] + self.assertEqual(cells[4:7], expected) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/runner/Cargo.toml b/tests/runner/Cargo.toml index 2941e85fa5a..35de1443aa8 100644 --- a/tests/runner/Cargo.toml +++ b/tests/runner/Cargo.toml @@ -13,7 +13,6 @@ clap.workspace = true client.workspace = true common-error.workspace = true common-query.workspace = true -common-recordbatch.workspace = true common-time.workspace = true datatypes = { workspace = true } flate2 = "1.0" diff --git a/tests/runner/src/client.rs b/tests/runner/src/client.rs index 163e827e53a..a595c877940 100644 --- a/tests/runner/src/client.rs +++ b/tests/runner/src/client.rs @@ -64,7 +64,11 @@ impl MultiProtocolClient { let grpc_client = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls(vec![grpc_server_addr]), + Client::with_query_and_control_managers( + Default::default(), + Default::default(), + vec![grpc_server_addr], + ), ); let pg_client = create_postgres_client(pg_server_addr).await; let mysql_client = create_mysql_client(mysql_server_addr).await; diff --git a/tests/runner/src/cmd/bare.rs b/tests/runner/src/cmd/bare.rs index 24ba756cee6..dd54d1d2d9b 100644 --- a/tests/runner/src/cmd/bare.rs +++ b/tests/runner/src/cmd/bare.rs @@ -112,7 +112,53 @@ pub struct BareCommand { enable_gc: bool, } +#[derive(Debug, PartialEq, Eq)] +struct Parallelism { + /// How many environments (standalone, distributed) run at the same time. + env: usize, + /// How many database instances run at the same time within one environment. + instance: usize, +} + impl BareCommand { + fn parallelism(&self) -> Parallelism { + let self_managed_kafka = + matches!(self.wal, Wal::Kafka) && self.kafka_wal_broker_endpoints.is_none(); + // External servers are shared by both environments, a filtered run is for + // debugging a single case, and each environment starts and tears down the + // self-managed Kafka cluster on its own. + if self.server_addr.server_addr.is_some() + || self.config.test_filter != ".*" + || self_managed_kafka + { + return Parallelism { + env: 1, + instance: 1, + }; + } + + // Instances of one environment would share the same metadata table and + // Kafka topics. Only distributed uses the kv backend, and the two + // environments use different Kafka topic prefixes, so the environments + // themselves can still run concurrently. + if self.setup_etcd + || !self.store_addrs.is_empty() + || self.setup_pg.is_some() + || self.setup_mysql.is_some() + || self.kafka_wal_broker_endpoints.is_some() + { + return Parallelism { + env: self.jobs, + instance: 1, + }; + } + + Parallelism { + env: self.jobs, + instance: self.jobs, + } + } + pub async fn run(mut self) { let temp_dir = tempfile::Builder::new() .prefix("sqlness") @@ -134,22 +180,11 @@ impl BareCommand { if self.jobs == 0 { self.jobs = num_cpus::get() / 2; } - - // normalize parallelism to 1 if any of the following conditions are met: - // Note: parallelism in pg and mysql is possible, but need configuration. - if self.server_addr.server_addr.is_some() - || self.setup_etcd - || self.setup_pg.is_some() - || self.setup_mysql.is_some() - || matches!(self.wal, Wal::Kafka) - || self.kafka_wal_broker_endpoints.is_some() - || self.config.test_filter != ".*" - { - self.jobs = 1; - println!( - "Normalizing parallelism to 1 due to server addresses, etcd/pg/mysql/kafka setup, or test filter usage" - ); - } + let parallelism = self.parallelism(); + println!( + "Environment parallelism: {}, instance parallelism per environment: {}", + parallelism.env, parallelism.instance + ); let config = ConfigBuilder::default() .case_dir(util::get_case_dir(self.config.case_dir)) @@ -158,8 +193,8 @@ impl BareCommand { .follow_links(true) .env_config_file(self.config.env_config_file) .interceptor_registry(interceptor_registry) - .parallelism(self.jobs) - .env_parallelism(self.jobs) + .parallelism(parallelism.instance) + .env_parallelism(parallelism.env) .build() .unwrap(); @@ -216,3 +251,61 @@ impl BareCommand { } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn parallelism(jobs: &str, args: &[&str]) -> Parallelism { + let cmd = BareCommand::try_parse_from(["bare", "-j", jobs].iter().chain(args)).unwrap(); + cmd.parallelism() + } + + #[test] + fn test_parallelism() { + let serial = Parallelism { + env: 1, + instance: 1, + }; + let env_only = Parallelism { + env: 4, + instance: 1, + }; + + assert_eq!( + parallelism("4", &[]), + Parallelism { + env: 4, + instance: 4 + } + ); + assert_eq!(parallelism("4", &["--setup-pg"]), env_only); + assert_eq!(parallelism("4", &["--setup-mysql"]), env_only); + assert_eq!(parallelism("4", &["--setup-etcd"]), env_only); + assert_eq!( + parallelism("4", &["--store-addrs", "127.0.0.1:2379"]), + env_only + ); + assert_eq!( + parallelism("4", &["-w", "kafka", "-k", "127.0.0.1:9092"]), + env_only + ); + assert_eq!(parallelism("4", &["-w", "kafka"]), serial); + assert_eq!(parallelism("4", &["-t", "basic"]), serial); + assert_eq!( + parallelism( + "4", + &[ + "-s", + "127.0.0.1:4001", + "-p", + "127.0.0.1:4003", + "-m", + "127.0.0.1:4002" + ] + ), + serial + ); + assert_eq!(parallelism("1", &["--setup-pg"]), serial); + } +} diff --git a/tests/runner/src/env/bare.rs b/tests/runner/src/env/bare.rs index dccd3673b16..43bbd987279 100644 --- a/tests/runner/src/env/bare.rs +++ b/tests/runner/src/env/bare.rs @@ -694,7 +694,7 @@ impl Env { "--bin", "greptime", "--features", - "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions", + "pg_kvbackend,mysql_kvbackend,ai_functions", ]) .output() .expect("Failed to start GreptimeDB"); diff --git a/tests/upgrade-compat/distributed/common b/tests/upgrade-compat/distributed/common deleted file mode 120000 index 2b0920287dc..00000000000 --- a/tests/upgrade-compat/distributed/common +++ /dev/null @@ -1 +0,0 @@ -../standalone/common \ No newline at end of file diff --git a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.result b/tests/upgrade-compat/standalone/common/table_engine_0_10_2.result deleted file mode 100644 index 046255a641f..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.result +++ /dev/null @@ -1,137 +0,0 @@ --- SQLNESS ARG version=v0.10.2 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -Affected Rows: 0 - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -Affected Rows: 0 - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -Affected Rows: 0 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -+---------------------+-----------------------------------------------------------+ -| Table | Create Table | -+---------------------+-----------------------------------------------------------+ -| mito_system_metrics | CREATE TABLE IF NOT EXISTS "mito_system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------------+-----------------------------------------------------------+ - -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "cpu_util" DOUBLE NULL, | -| | "host" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host") | -| | ) | -| | | -| | ENGINE=metric | -| | WITH( | -| | on_physical_table = 'phy' | -| | ) | -+----------------+-----------------------------------------------------------+ - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -Affected Rows: 3 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -Affected Rows: 3 - -SELECT * FROM mito_system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host3 | idc_a | 90.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host4 | idc_a | 70.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host5 | idc_a | 60.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - -SELECT * FROM system_metrics; - -+----------+-------+-------------------------+ -| cpu_util | host | ts | -+----------+-------+-------------------------+ -| 80.0 | host2 | 2022-11-03T03:39:57.450 | -| 70.0 | host4 | 2022-11-03T03:39:57.450 | -| 60.0 | host5 | 2022-11-03T03:39:57.450 | -| 90.0 | host3 | 2022-11-03T03:39:57.450 | -| 50.0 | host1 | 2022-11-03T03:39:57.450 | -+----------+-------+-------------------------+ - -DROP TABLE mito_system_metrics; - -Affected Rows: 0 - -DROP TABLE system_metrics; - -Affected Rows: 0 - -DROP TABLE phy; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.sql b/tests/upgrade-compat/standalone/common/table_engine_0_10_2.sql deleted file mode 100644 index 1907533b159..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.sql +++ /dev/null @@ -1,60 +0,0 @@ --- SQLNESS ARG version=v0.10.2 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -SHOW CREATE TABLE system_metrics; - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -SELECT * FROM mito_system_metrics; - -SELECT * FROM system_metrics; - -DROP TABLE mito_system_metrics; - -DROP TABLE system_metrics; - -DROP TABLE phy; diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.result b/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.result deleted file mode 100644 index 7ce230a6881..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.result +++ /dev/null @@ -1,137 +0,0 @@ --- SQLNESS ARG version=v0.11.0 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -Affected Rows: 0 - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -Affected Rows: 0 - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -Affected Rows: 0 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -+---------------------+-----------------------------------------------------------+ -| Table | Create Table | -+---------------------+-----------------------------------------------------------+ -| mito_system_metrics | CREATE TABLE IF NOT EXISTS "mito_system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------------+-----------------------------------------------------------+ - -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "cpu_util" DOUBLE NULL, | -| | "host" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host") | -| | ) | -| | | -| | ENGINE=metric | -| | WITH( | -| | on_physical_table = 'phy' | -| | ) | -+----------------+-----------------------------------------------------------+ - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -Affected Rows: 3 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -Affected Rows: 3 - -SELECT * FROM mito_system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host3 | idc_a | 90.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host4 | idc_a | 70.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host5 | idc_a | 60.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - -SELECT * FROM system_metrics; - -+----------+-------+-------------------------+ -| cpu_util | host | ts | -+----------+-------+-------------------------+ -| 80.0 | host2 | 2022-11-03T03:39:57.450 | -| 70.0 | host4 | 2022-11-03T03:39:57.450 | -| 60.0 | host5 | 2022-11-03T03:39:57.450 | -| 90.0 | host3 | 2022-11-03T03:39:57.450 | -| 50.0 | host1 | 2022-11-03T03:39:57.450 | -+----------+-------+-------------------------+ - -DROP TABLE mito_system_metrics; - -Affected Rows: 0 - -DROP TABLE system_metrics; - -Affected Rows: 0 - -DROP TABLE phy; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.sql b/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.sql deleted file mode 100644 index 963170fdf58..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.sql +++ /dev/null @@ -1,60 +0,0 @@ --- SQLNESS ARG version=v0.11.0 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -SHOW CREATE TABLE system_metrics; - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -SELECT * FROM mito_system_metrics; - -SELECT * FROM system_metrics; - -DROP TABLE mito_system_metrics; - -DROP TABLE system_metrics; - -DROP TABLE phy; diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.result b/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.result deleted file mode 100644 index 41b81f01c08..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.result +++ /dev/null @@ -1,137 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -Affected Rows: 0 - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -Affected Rows: 0 - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -Affected Rows: 0 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -+---------------------+-----------------------------------------------------------+ -| Table | Create Table | -+---------------------+-----------------------------------------------------------+ -| mito_system_metrics | CREATE TABLE IF NOT EXISTS "mito_system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------------+-----------------------------------------------------------+ - -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "cpu_util" DOUBLE NULL, | -| | "host" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host") | -| | ) | -| | | -| | ENGINE=metric | -| | WITH( | -| | on_physical_table = 'phy' | -| | ) | -+----------------+-----------------------------------------------------------+ - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -Affected Rows: 3 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -Affected Rows: 3 - -SELECT * FROM mito_system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host3 | idc_a | 90.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host4 | idc_a | 70.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host5 | idc_a | 60.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - -SELECT * FROM system_metrics; - -+----------+-------+-------------------------+ -| cpu_util | host | ts | -+----------+-------+-------------------------+ -| 80.0 | host2 | 2022-11-03T03:39:57.450 | -| 70.0 | host4 | 2022-11-03T03:39:57.450 | -| 60.0 | host5 | 2022-11-03T03:39:57.450 | -| 90.0 | host3 | 2022-11-03T03:39:57.450 | -| 50.0 | host1 | 2022-11-03T03:39:57.450 | -+----------+-------+-------------------------+ - -DROP TABLE mito_system_metrics; - -Affected Rows: 0 - -DROP TABLE system_metrics; - -Affected Rows: 0 - -DROP TABLE phy; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.sql b/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.sql deleted file mode 100644 index 9908085213e..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.sql +++ /dev/null @@ -1,60 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -SHOW CREATE TABLE system_metrics; - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -SELECT * FROM mito_system_metrics; - -SELECT * FROM system_metrics; - -DROP TABLE mito_system_metrics; - -DROP TABLE system_metrics; - -DROP TABLE phy; diff --git a/tests/upgrade-compat/standalone/common/test_alter_drop_column.result b/tests/upgrade-compat/standalone/common/test_alter_drop_column.result deleted file mode 100644 index 605f8711642..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_drop_column.result +++ /dev/null @@ -1,62 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_drop_test ( - id INT, - val_to_drop STRING, - val_to_keep STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -Affected Rows: 0 - -INSERT INTO alter_drop_test(id, val_to_drop, val_to_keep, ts) VALUES (1, 'a', 'x', 1672531200000); - -Affected Rows: 1 - --- SQLNESS ARG version=latest -ALTER TABLE alter_drop_test DROP COLUMN val_to_drop; - -Affected Rows: 0 - -SHOW CREATE TABLE alter_drop_test; - -+-----------------+------------------------------------------------+ -| Table | Create Table | -+-----------------+------------------------------------------------+ -| alter_drop_test | CREATE TABLE IF NOT EXISTS "alter_drop_test" ( | -| | "id" INT NULL, | -| | "val_to_keep" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("id") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+-----------------+------------------------------------------------+ - -SELECT * FROM alter_drop_test ORDER BY id; - -+----+-------------+---------------------+ -| id | val_to_keep | ts | -+----+-------------+---------------------+ -| 1 | x | 2023-01-01T00:00:00 | -+----+-------------+---------------------+ - -INSERT INTO alter_drop_test(id, val_to_keep, ts) VALUES (2, 'y', 1672531201000); - -Affected Rows: 1 - -SELECT * FROM alter_drop_test ORDER BY id; - -+----+-------------+---------------------+ -| id | val_to_keep | ts | -+----+-------------+---------------------+ -| 1 | x | 2023-01-01T00:00:00 | -| 2 | y | 2023-01-01T00:00:01 | -+----+-------------+---------------------+ - -DROP TABLE alter_drop_test; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_alter_drop_column.sql b/tests/upgrade-compat/standalone/common/test_alter_drop_column.sql deleted file mode 100644 index df52affe522..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_drop_column.sql +++ /dev/null @@ -1,23 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_drop_test ( - id INT, - val_to_drop STRING, - val_to_keep STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -INSERT INTO alter_drop_test(id, val_to_drop, val_to_keep, ts) VALUES (1, 'a', 'x', 1672531200000); - --- SQLNESS ARG version=latest -ALTER TABLE alter_drop_test DROP COLUMN val_to_drop; - -SHOW CREATE TABLE alter_drop_test; - -SELECT * FROM alter_drop_test ORDER BY id; - -INSERT INTO alter_drop_test(id, val_to_keep, ts) VALUES (2, 'y', 1672531201000); - -SELECT * FROM alter_drop_test ORDER BY id; - -DROP TABLE alter_drop_test; diff --git a/tests/upgrade-compat/standalone/common/test_alter_table.result b/tests/upgrade-compat/standalone/common/test_alter_table.result deleted file mode 100644 index df9bbe1a91c..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_table.result +++ /dev/null @@ -1,54 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_test ( - id INT, - val STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -Affected Rows: 0 - -INSERT INTO alter_test(id, val, ts) VALUES (1, 'a', 1672531200000); - -Affected Rows: 1 - --- SQLNESS ARG version=latest -ALTER TABLE alter_test ADD COLUMN new_col INT; - -Affected Rows: 0 - -SHOW CREATE TABLE alter_test; - -+------------+-------------------------------------------+ -| Table | Create Table | -+------------+-------------------------------------------+ -| alter_test | CREATE TABLE IF NOT EXISTS "alter_test" ( | -| | "id" INT NULL, | -| | "val" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "new_col" INT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("id") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+------------+-------------------------------------------+ - -INSERT INTO alter_test(id, val, ts, new_col) VALUES (2, 'b', 1672531201000, 100); - -Affected Rows: 1 - -SELECT * FROM alter_test ORDER BY id; - -+----+-----+---------------------+---------+ -| id | val | ts | new_col | -+----+-----+---------------------+---------+ -| 1 | a | 2023-01-01T00:00:00 | | -| 2 | b | 2023-01-01T00:00:01 | 100 | -+----+-----+---------------------+---------+ - -DROP TABLE alter_test; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_alter_table.sql b/tests/upgrade-compat/standalone/common/test_alter_table.sql deleted file mode 100644 index 41b1555d8a8..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_table.sql +++ /dev/null @@ -1,20 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_test ( - id INT, - val STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -INSERT INTO alter_test(id, val, ts) VALUES (1, 'a', 1672531200000); - --- SQLNESS ARG version=latest -ALTER TABLE alter_test ADD COLUMN new_col INT; - -SHOW CREATE TABLE alter_test; - -INSERT INTO alter_test(id, val, ts, new_col) VALUES (2, 'b', 1672531201000, 100); - -SELECT * FROM alter_test ORDER BY id; - -DROP TABLE alter_test; diff --git a/tests/upgrade-compat/standalone/common/test_simple.result b/tests/upgrade-compat/standalone/common/test_simple.result deleted file mode 100644 index ff2c340598a..00000000000 --- a/tests/upgrade-compat/standalone/common/test_simple.result +++ /dev/null @@ -1,47 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -); - -Affected Rows: 0 - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+----------------+-----------------------------------------------------------+ - -DROP TABLE system_metrics; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_simple.sql b/tests/upgrade-compat/standalone/common/test_simple.sql deleted file mode 100644 index 0f8daa0985a..00000000000 --- a/tests/upgrade-compat/standalone/common/test_simple.sql +++ /dev/null @@ -1,22 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -); - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE system_metrics; - -DROP TABLE system_metrics; diff --git a/tests/upgrade-compat/standalone/common/test_ttl.result b/tests/upgrade-compat/standalone/common/test_ttl.result deleted file mode 100644 index c368d13e4c1..00000000000 --- a/tests/upgrade-compat/standalone/common/test_ttl.result +++ /dev/null @@ -1,153 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE test_ttl_0s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '0 second'); - -Affected Rows: 0 - -CREATE TABLE test_ttl_1s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '1 second'); - -Affected Rows: 0 - -CREATE TABLE test_ttl_none(ts TIMESTAMP TIME INDEX, val INT); - -Affected Rows: 0 - -CREATE DATABASE ttl_db_1s WITH (ttl = '1 second'); - -Affected Rows: 1 - -CREATE DATABASE ttl_db_0s WITH (ttl = '0 second'); - -Affected Rows: 1 - -CREATE DATABASE ttl_db_none; - -Affected Rows: 1 - --- SQLNESS ARG version=latest -SHOW TABLES; - -+------------------+ -| Tables_in_public | -+------------------+ -| numbers | -| test_ttl_0s | -| test_ttl_1s | -| test_ttl_none | -+------------------+ - -SHOW CREATE TABLE test_ttl_1s; - -+-------------+--------------------------------------------+ -| Table | Create Table | -+-------------+--------------------------------------------+ -| test_ttl_1s | CREATE TABLE IF NOT EXISTS "test_ttl_1s" ( | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "val" INT NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | ttl = '1s' | -| | ) | -+-------------+--------------------------------------------+ - -SHOW CREATE TABLE test_ttl_0s; - -+-------------+--------------------------------------------+ -| Table | Create Table | -+-------------+--------------------------------------------+ -| test_ttl_0s | CREATE TABLE IF NOT EXISTS "test_ttl_0s" ( | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "val" INT NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | ttl = '0s' | -| | ) | -+-------------+--------------------------------------------+ - -SHOW CREATE TABLE test_ttl_none; - -+---------------+----------------------------------------------+ -| Table | Create Table | -+---------------+----------------------------------------------+ -| test_ttl_none | CREATE TABLE IF NOT EXISTS "test_ttl_none" ( | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "val" INT NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------+----------------------------------------------+ - -DROP TABLE test_ttl_1s; - -Affected Rows: 0 - -DROP TABLE test_ttl_0s; - -Affected Rows: 0 - -DROP TABLE test_ttl_none; - -Affected Rows: 0 - -SHOW DATABASES; - -+--------------------+ -| Database | -+--------------------+ -| greptime_private | -| information_schema | -| public | -| ttl_db_0s | -| ttl_db_1s | -| ttl_db_none | -+--------------------+ - -SHOW CREATE DATABASE ttl_db_1s; - -+-----------+-----------------------------------------+ -| Database | Create Database | -+-----------+-----------------------------------------+ -| ttl_db_1s | CREATE DATABASE IF NOT EXISTS ttl_db_1s | -| | WITH( | -| | ttl = '1s' | -| | ) | -+-----------+-----------------------------------------+ - -SHOW CREATE DATABASE ttl_db_0s; - -+-----------+-----------------------------------------+ -| Database | Create Database | -+-----------+-----------------------------------------+ -| ttl_db_0s | CREATE DATABASE IF NOT EXISTS ttl_db_0s | -| | WITH( | -| | ttl = '0s' | -| | ) | -+-----------+-----------------------------------------+ - -SHOW CREATE DATABASE ttl_db_none; - -+-------------+-------------------------------------------+ -| Database | Create Database | -+-------------+-------------------------------------------+ -| ttl_db_none | CREATE DATABASE IF NOT EXISTS ttl_db_none | -+-------------+-------------------------------------------+ - -DROP DATABASE ttl_db_1s; - -Affected Rows: 0 - -DROP DATABASE ttl_db_0s; - -Affected Rows: 0 - -DROP DATABASE ttl_db_none; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_ttl.sql b/tests/upgrade-compat/standalone/common/test_ttl.sql deleted file mode 100644 index 3462fd22444..00000000000 --- a/tests/upgrade-compat/standalone/common/test_ttl.sql +++ /dev/null @@ -1,42 +0,0 @@ - --- SQLNESS ARG version=v0.9.5 -CREATE TABLE test_ttl_0s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '0 second'); - -CREATE TABLE test_ttl_1s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '1 second'); - -CREATE TABLE test_ttl_none(ts TIMESTAMP TIME INDEX, val INT); - -CREATE DATABASE ttl_db_1s WITH (ttl = '1 second'); - -CREATE DATABASE ttl_db_0s WITH (ttl = '0 second'); - -CREATE DATABASE ttl_db_none; - --- SQLNESS ARG version=latest -SHOW TABLES; - -SHOW CREATE TABLE test_ttl_1s; - -SHOW CREATE TABLE test_ttl_0s; - -SHOW CREATE TABLE test_ttl_none; - -DROP TABLE test_ttl_1s; - -DROP TABLE test_ttl_0s; - -DROP TABLE test_ttl_none; - -SHOW DATABASES; - -SHOW CREATE DATABASE ttl_db_1s; - -SHOW CREATE DATABASE ttl_db_0s; - -SHOW CREATE DATABASE ttl_db_none; - -DROP DATABASE ttl_db_1s; - -DROP DATABASE ttl_db_0s; - -DROP DATABASE ttl_db_none; diff --git a/tests/upgrade-compat/standalone/common/test_update.result b/tests/upgrade-compat/standalone/common/test_update.result deleted file mode 100644 index 872878cb152..00000000000 --- a/tests/upgrade-compat/standalone/common/test_update.result +++ /dev/null @@ -1,32 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE update_test ( - id INT, - val INT, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -Affected Rows: 0 - -INSERT INTO update_test(id, val, ts) VALUES (1, 10, 1672531200000), (2, 20, 1672531201000); - -Affected Rows: 2 - --- SQLNESS ARG version=latest -UPDATE update_test SET val = 100 WHERE id = 1; - -Error: 1001(Unsupported), SQL statement is not supported, keyword: UPDATE - -SELECT * FROM update_test ORDER BY id; - -+----+-----+---------------------+ -| id | val | ts | -+----+-----+---------------------+ -| 1 | 10 | 2023-01-01T00:00:00 | -| 2 | 20 | 2023-01-01T00:00:01 | -+----+-----+---------------------+ - -DROP TABLE update_test; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_update.sql b/tests/upgrade-compat/standalone/common/test_update.sql deleted file mode 100644 index 33eedcd2061..00000000000 --- a/tests/upgrade-compat/standalone/common/test_update.sql +++ /dev/null @@ -1,16 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE update_test ( - id INT, - val INT, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -INSERT INTO update_test(id, val, ts) VALUES (1, 10, 1672531200000), (2, 20, 1672531201000); - --- SQLNESS ARG version=latest -UPDATE update_test SET val = 100 WHERE id = 1; - -SELECT * FROM update_test ORDER BY id; - -DROP TABLE update_test;