rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 15:02:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into fix/windows-fs-paths

Browse Source
This commit is contained in:
Weny Xu
2026-09-23 10:56:39 +08:00
committed by GitHub
parent 9ed4c008ff 0f118bc5ba
commit 4cdd6d757f
131 changed files with 8260 additions and 1979 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/actions/build-linux-artifacts/action.yml
+23
View File
@@ -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 \
.github/scripts/run-fuzz-targets-test.sh
+18 -3
View File
@@ -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'
.github/scripts/run-fuzz-targets.sh
+2 -2
View File
@@ -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
.github/workflows/dev-build.yml
+8
View File
@@ -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
.github/workflows/integration.yml
+15 -2
View File
@@ -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 }}
@@ -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 }}
.github/workflows/nightly-build.yml
+8
View File
@@ -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
.github/workflows/nightly-ci.yml
+52 -1
View File
@@ -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 vector_index -F ai_functions
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:
.github/workflows/release.yml
+8
View File
@@ -204,6 +204,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
@@ -227,6 +231,10 @@ 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
.github/workflows/rust.yml
+16 -18
View File
@@ -248,15 +248,14 @@ jobs:
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
@@ -305,15 +304,14 @@ jobs:
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
Cargo.lock
Generated
+2 -1
View File
@@ -6130,7 +6130,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=3d232f406fad574e8c630fb1c304fe2ebb013c9c#3d232f406fad574e8c630fb1c304fe2ebb013c9c"
dependencies = [
"prost 0.14.1",
"prost-types 0.14.1",
@@ -14457,6 +14457,7 @@ version = "1.3.0-alpha.1"
dependencies = [
"api",
"aquamarine",
"arrow-schema 59.2.0",
"async-stream",
"async-trait",
"bytes",
Cargo.toml
+1 -1
View File
@@ -161,7 +161,7 @@ fs2 = "0.4"
fst = "0.4.7"
futures = "0.3"
futures-util = "0.3"
greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "0c4beffaa2b3968034915867593a717c6ab0c6b4" }
greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "3d232f406fad574e8c630fb1c304fe2ebb013c9c" }
hex = "0.4"
hostname = "0.4.0"
http = "1"
Makefile
+5 -4
View File
@@ -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.
@@ -264,6 +264,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}
config/config.md
+20 -18
View File
@@ -34,7 +34,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.<br/>When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is<br/>raised to at least the largest active flush interval plus 1 second. The intervals come from<br/>`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.<br/>When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is<br/>raised to at least the largest active flush interval plus 1 second. The intervals come from<br/>`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.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>This allows browser to access http APIs without CORS restrictions |
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
@@ -75,13 +75,20 @@
| `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.<br/>Available values: "last_non_null", "last_row". |
| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.<br/>Legacy Prometheus batching settings under prom_store remain supported.<br/>HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.<br/>Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.<br/>Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.<br/>Effective shared Prom settings take precedence; existing prom_store settings remain compatible. |
| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in HTTP ingestion protocols.<br/>PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge<br/>queue admission without waiting for storage; later failures cannot be returned to the client.<br/>Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.<br/>OTLP logs, traces and ordinary metrics use this batcher. |
| `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.<br/>Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.<br/>Enable independently with protocols and a nonzero flush interval.<br/>Omitted fields use independent defaults, not parent settings.<br/>Empty protocols or a zero interval disables logical batching without fallback.<br/>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. |
@@ -94,12 +101,6 @@
| `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.<br/>Available options:<br/>- strict: deny invalid UTF-8 strings (default).<br/>- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).<br/>- 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.<br/>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.<br/>- `raft_engine`: the wal is stored in the local file system by raft-engine.<br/>- `kafka`: it's remote wal that data is stored in Kafka.<br/>- `experimental_object_store`: the wal is stored as objects in an object store.<br/>**Notes: experimental and not supported yet.** |
| `wal.dir` | String | Unset | The directory to store the WAL files.<br/>**It's only used when the provider is `raft_engine`**. |
@@ -289,7 +290,7 @@
| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.<br/>Defaults to max(num_cpus / 2, 2). |
| `http` | -- | -- | The HTTP server options. |
| `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. |
| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.<br/>When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is<br/>raised to at least the largest active flush interval plus 1 second. The intervals come from<br/>`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.<br/>When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is<br/>raised to at least the largest active flush interval plus 1 second. The intervals come from<br/>`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.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>This allows browser to access http APIs without CORS restrictions |
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
@@ -342,13 +343,20 @@
| `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.<br/>Available values: "last_non_null", "last_row". |
| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.<br/>Legacy Prometheus batching settings under prom_store remain supported.<br/>HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.<br/>Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.<br/>Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.<br/>Effective shared Prom settings take precedence; existing prom_store settings remain compatible. |
| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in HTTP ingestion protocols.<br/>PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge<br/>queue admission without waiting for storage; later failures cannot be returned to the client.<br/>Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.<br/>OTLP logs, traces and ordinary metrics use this batcher. |
| `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.<br/>Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.<br/>Enable independently with protocols and a nonzero flush interval.<br/>Omitted fields use independent defaults, not parent settings.<br/>Empty protocols or a zero interval disables logical batching without fallback.<br/>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. |
@@ -361,12 +369,6 @@
| `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.<br/>Available options:<br/>- strict: deny invalid UTF-8 strings (default).<br/>- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).<br/>- 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.<br/>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. |
config/frontend.example.toml
+30 -22
View File
@@ -56,9 +56,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`.
@@ -231,12 +231,11 @@ 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 HTTP 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.
[pending_rows_batcher]
# protocols = [
# "influxdb",
@@ -256,12 +255,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.
@@ -293,19 +313,7 @@ with_metric_engine = true
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]
config/standalone.example.toml
+30 -22
View File
@@ -81,9 +81,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`.
@@ -210,12 +210,11 @@ 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 HTTP 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.
[pending_rows_batcher]
# protocols = [
# "influxdb",
@@ -235,12 +234,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.
@@ -272,19 +292,7 @@ with_metric_engine = true
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]
docs/how-to/ai-functions.md
+247
View File
@@ -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='<your TypeSafe 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: <https://docs.typesafe.ai/api>
## 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
```
docs/how-to/jev.md
-119
View File
@@ -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='<your TypeSafe 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: <https://docs.typesafe.ai/api>
## 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
```
src/cmd/src/bin/query_perf_fixture/case.rs
+1 -1
View File
@@ -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
src/cmd/tests/load_config_test.rs
+9
View File
@@ -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;
@@ -206,6 +207,10 @@ fn test_load_frontend_example_config() {
);
let expected = GreptimeOptions::<FrontendOptions> {
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,
@@ -389,6 +394,10 @@ fn test_load_standalone_example_config() {
);
let expected = GreptimeOptions::<StandaloneOptions> {
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,
src/common/event-recorder/src/recorder.rs
+22 -8
View File
@@ -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<dyn EventHandler>) -> 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<dyn EventHandler>,
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(
src/common/function/src/function_registry.rs
+14 -8
View File
@@ -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<Arc<FunctionRegistry>> = 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]
src/common/function/src/scalars.rs
+2 -2
View File
@@ -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;
src/common/function/src/scalars/ai.rs
+454
View File
@@ -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::<Noul>::register(registry);
AiFunction::<Choice>::register(registry);
AiFunction::<Score>::register(registry);
}
/// Shared asynchronous execution for AI functions using the Jev backend.
#[derive(PartialEq, Eq, Hash)]
struct AiFunction<Q> {
signature: Signature,
question: PhantomData<Q>,
enabled: bool,
api_key: Option<String>,
endpoint: String,
model: String,
}
struct AiRequest<'a, Q: AiQuestion> {
text: &'a str,
prompt: &'a str,
criteria: Arc<Q::Criteria>,
}
/// 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<Self> {
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<Q: AiQuestion> AiFunction<Q> {
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<ScalarValue> {
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<Vec<Option<AiRequest<'a, Q>>>> {
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::<Result<Vec<_>>>()?;
let scalar_criteria = args[2..]
.iter()
.all(|arg| matches!(arg, StringArgument::Scalar(_)));
let mut shared_criteria: Option<Arc<Q::Criteria>> = None;
// Validate all non-null rows before making any billable requests.
(0..number_rows)
.map(|row| {
let values: Option<Vec<_>> = 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<Client> {
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<Q: AiQuestion> Default for AiFunction<Q> {
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<Q: AiQuestion> fmt::Debug for AiFunction<Q> {
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<Q: AiQuestion> ScalarUDFImpl for AiFunction<Q> {
fn name(&self) -> &str {
Q::NAME
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(Q::return_type())
}
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
not_impl_err!("{} can only be called from async contexts", Q::NAME)
}
}
#[async_trait]
impl<Q: AiQuestion> AsyncScalarUDFImpl for AiFunction<Q> {
async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
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<ScalarValue> = 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<Self::Criteria>;
fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result<ScalarValue>;
}
#[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<ScalarValue> {
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<String, Value>;
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<Self::Criteria> {
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<ScalarValue> {
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<Value>;
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<Self::Criteria> {
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<ScalarValue> {
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<Vec<f64>> {
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::<Result<Vec<_>>>()?;
// Allow small rounding differences without renormalizing the provider's distribution.
if (probabilities.iter().sum::<f64>() - 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<f64> {
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}]")
})
}
src/common/function/src/scalars/ai/tests.rs
+766
View File
@@ -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<Q: AiQuestion>(&self) -> AiFunction<Q> {
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<Q: AiQuestion>(function: AiFunction<Q>) -> 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<Q: AiQuestion>(criteria: &str) {
let function = AiFunction::<Q>::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::<Choice>(r#"{"billing":null,"technical":"errors"}"#);
assert_shared_criteria::<Score>(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<Value>| 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::<Noul>());
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::<Noul> {
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::<Noul> {
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::<Noul>());
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<Value>| 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::<Choice>());
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<Value>| 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<String, Value> = (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::<Score>());
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<Value>| {
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::<Score>());
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<f64> = 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::<Score>());
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::<Choice> {
enabled: true,
api_key: None,
..Default::default()
});
let score_ctx = context(AiFunction::<Score> {
enabled: true,
api_key: None,
..Default::default()
});
let too_many_choices: Map<String, Value> = (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::<Choice>());
ctx.register_udf(
AsyncScalarUDF::new(Arc::new(server.function::<Score>())).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);
}
src/common/function/src/scalars/jev.rs
-209
View File
@@ -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<String>,
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<f64> {
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<Client> {
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<DataType> {
Ok(DataType::Boolean)
}
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
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<ColumnarValue> {
let arrays = args
.args
.into_iter()
.map(|arg| arg.into_array(args.number_rows))
.collect::<Result<Vec<_>>>()?;
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<Option<bool>> = stream::iter(requests).buffered(8).try_collect().await?;
Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(matches))))
}
}
src/common/function/src/scalars/jev/tests.rs
-250
View File
@@ -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<Value>| 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);
}
src/common/grpc-expr/src/alter.rs
+46 -5
View File
@@ -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<JsonSettings> {
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::<Result<Vec<_>>>()?;
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(),
},
src/common/grpc-expr/src/error.rs
+8
View File
@@ -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 { .. } => {
src/common/meta/src/ddl/alter_table/executor.rs
+1
View File
@@ -333,6 +333,7 @@ fn build_new_table_info(
}
AlterKind::DropColumns { .. }
| AlterKind::ModifyColumnTypes { .. }
| AlterKind::SetJsonSettings { .. }
| AlterKind::SetTableOptions { .. }
| AlterKind::UnsetTableOptions { .. }
| AlterKind::SetAnnotations { .. }
src/common/meta/src/ddl/alter_table/region_request.rs
+1
View File
@@ -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
src/common/meta/src/ddl/event/table.rs
+1
View File
@@ -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"),
src/datatypes/src/data_type.rs
+3 -1
View File
@@ -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(_)
src/datatypes/src/extension/json.rs
+53 -2
View File
@@ -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<Metadata> {
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<T: AsRef<Field>>(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)
src/datatypes/src/json.rs
+336 -48
View File
@@ -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::<Vec<_>>();
let mut other_hints = other_hints.iter().collect::<Vec<_>>();
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<Json> {
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<Value> {
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<Value> {
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)))
}
}
@@ -272,36 +316,41 @@ fn with_key_context<T>(
}
/// Main encoding function with key path tracking
fn encode_json_with_context(json: Json, context: &mut JsonContext) -> Result<JsonValue> {
fn encode_json_with_context(
json: Json,
context: &mut JsonContext,
policy: TypeHintMismatchPolicy,
) -> Result<JsonValue> {
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<String, Json>,
context: &mut JsonContext<'a>,
policy: TypeHintMismatchPolicy,
) -> Result<JsonValue> {
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 +358,14 @@ fn encode_json_object_with_context<'a>(
fn fill_missing_type_hints(
object: &mut BTreeMap<String, JsonVariant>,
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 +377,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 +392,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 +409,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 +420,152 @@ fn encode_json_value_with_hint(
json: Json,
hint: &JsonTypeHint,
context: &mut JsonContext,
policy: TypeHintMismatchPolicy,
) -> Result<JsonValue> {
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::<Vec<_>>();
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::<Vec<_>>();
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<Json>,
context: &mut JsonContext<'a>,
policy: TypeHintMismatchPolicy,
) -> Result<JsonValue> {
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 +600,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<JsonValue> {
fn encode_json_value_with_context(
json: Json,
context: &mut JsonContext,
policy: TypeHintMismatchPolicy,
) -> Result<JsonValue> {
if context.path.len() >= JSON2_MAX_STRUCTURED_DEPTH
&& matches!(&json, Json::Object(_) | Json::Array(_))
{
@@ -487,8 +633,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 +743,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<dyn std::error::Error>> {
@@ -755,6 +945,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();
src/datatypes/src/types/struct_type.rs
+19
View File
@@ -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<Self> {
let struct_fields = fields
.iter()
.map(|field| {
Ok(StructField::new(
field.name().clone(),
ConcreteDataType::try_from(field.data_type())?,
field.is_nullable(),
))
})
.collect::<Result<Vec<_>>>()?;
Ok(StructType {
fields: Arc::new(struct_fields),
})
}
pub fn fields(&self) -> Arc<Vec<StructField>> {
self.fields.clone()
}
src/datatypes/src/value.rs
+94 -5
View File
@@ -1152,12 +1152,13 @@ impl TryFrom<ScalarValue> 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<ScalarValue> 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<ScalarValue> for Value {
| ScalarValue::LargeListView(_)
| ScalarValue::Union(_, _, _)
| ScalarValue::Float16(_)
| ScalarValue::Utf8View(_)
| ScalarValue::BinaryView(_)
| ScalarValue::Map(_)
| ScalarValue::Date64(_)
| ScalarValue::RunEndEncoded(_, _, _) => {
@@ -1259,6 +1266,21 @@ impl TryFrom<ScalarValue> 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<Value> {
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<ValueRef<'_>> 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]
src/datatypes/src/vectors/json/array.rs
+23 -82
View File
@@ -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<ArrayRef> {
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<ArrayRef> {
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<Value>, to_type: &DataType) -> Result<ArrayRe
let concrete_type = ConcreteDataType::from_arrow_type(to_type);
let mut builder = concrete_type.create_mutable_vector(values.len());
for value in values {
let value = project_json_value_to_type(value, &concrete_type)?;
let value = coerce_json_value_to_type(value, &concrete_type);
builder.try_push_value_ref(&value.as_value_ref())?;
}
Ok(builder.to_vector().to_arrow_array())
}
fn project_json_value_to_type(value: Value, to_type: &ConcreteDataType) -> Result<GreptimeValue> {
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::<Result<Vec<_>>>()?;
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::<Result<Vec<_>>>()?;
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 }
src/frontend/AGENTS.md
+5
View File
@@ -51,6 +51,11 @@ remote datanodes via `operator`/`client`.
- Internal gRPC listeners mark requests with `Channel::Internal` in middleware
(`server.rs`), including requests handled by Enterprise Flight wrappers.
- **Logical-table batching** (`instance/logical_batcher.rs`): `Services` initializes
one shared batcher for opted-in HTTP Prom and nonlegacy OTLP metric-engine
writes. OTLP checks operator eligibility 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
src/frontend/src/frontend.rs
+44 -5
View File
@@ -67,7 +67,7 @@ pub struct FrontendOptions {
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 +131,7 @@ impl Configurable for FrontendOptions {
"meta_client.metasrv_addrs",
"event_recorder.event_types",
"pending_rows_batcher.protocols",
"pending_rows_batcher.logical_table.protocols",
])
}
}
@@ -225,6 +226,38 @@ mod tests {
type GrpcStream<T> =
Pin<Box<dyn Stream<Item = std::result::Result<T, Status>> + 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![
servers::http::BatchingProtocol::Prom,
servers::http::BatchingProtocol::Otlp
]
);
},
);
}
#[test]
fn test_batcher_protocols_from_env() {
temp_env::with_vars(
@@ -236,7 +269,7 @@ mod tests {
let options =
FrontendOptions::load_layered_options(None, "FRONTEND_BATCHER_TEST").unwrap();
assert_eq!(
options.pending_rows_batcher.protocols,
options.pending_rows_batcher.table.protocols,
vec![
servers::http::BatchingProtocol::Influxdb,
servers::http::BatchingProtocol::HttpSql
@@ -252,6 +285,7 @@ mod tests {
assert!(
!defaults
.pending_rows_batcher
.table
.pending_rows_batching_enabled()
);
let options: FrontendOptions = toml::from_str(
@@ -263,9 +297,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);
src/frontend/src/instance.rs
+4 -1
View File
@@ -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;
@@ -78,6 +79,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 +132,7 @@ pub struct Instance {
query_engine: QueryEngineRef,
plugins: Plugins,
inserter: InserterRef,
logical_batcher: Arc<OnceLock<Option<Arc<LogicalTablePendingRowsBatcher>>>>,
deleter: DeleterRef,
table_metadata_manager: TableMetadataManagerRef,
event_recorder: EventRecorderRef,
src/frontend/src/instance/builder.rs
+25 -24
View File
@@ -61,7 +61,7 @@ use crate::heartbeat::frontend_peer_addr;
use crate::instance::Instance;
use crate::instance::entity_graph::EntityGraphProviderImpl;
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 +237,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<Arc<dyn PendingRowsBatcher>> {
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<dyn PendingRowsBatcher>)
};
let build_batcher = |options: &BatcherOptions| -> Option<Arc<dyn PendingRowsBatcher>> {
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<dyn PendingRowsBatcher>)
};
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 +381,12 @@ 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);
Ok(Instance {
logical_batcher: Default::default(),
frontend_peer_addr,
experimental_metric_export: self.options.experimental_metric_export,
catalog_manager: self.catalog_manager,
src/frontend/src/instance/logical_batcher.rs
+102
View File
@@ -0,0 +1,102 @@
// 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::logical_table::{LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer};
use servers::error::{BatcherChannelClosedSnafu, Result};
use servers::http::BatchingProtocol;
use session::context::QueryContextRef;
use snafu::OptionExt;
use crate::frontend::FrontendOptions;
use crate::instance::Instance;
impl Instance {
pub(crate) fn init_logical_batcher(self: &Arc<Self>, options: &FrontendOptions) {
self.logical_batcher.get_or_init(|| {
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 None;
}
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<LogicalTablePendingRowsBatcher>> {
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<Instance>);
#[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
}
}
src/frontend/src/instance/otlp.rs
+45 -12
View File
@@ -27,10 +27,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};
@@ -165,27 +166,59 @@ impl OpenTelemetryProtocolHandler for Instance {
Arc::new(c)
};
let physical_table = ctx
.extension(PHYSICAL_TABLE_PARAM)
.unwrap_or(GREPTIME_PHYSICAL_TABLE)
.to_string();
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.
let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine {
let output = 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?;
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)
.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)
}?;
.context(error::ExecuteGrpcQuerySnafu)?
};
outcome.write_cost = output.meta.cost;
// Derived enrichment, written after the metric data is committed:
// failing here would make the client retry data the server already
// accepted, so every failure degrades to a warning instead.
// 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,
src/frontend/src/server.rs
+126 -45
View File
@@ -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::pending_rows_batch_sync_enabled;
use servers::error::Error as ServerError;
use servers::grpc::builder::GrpcServerBuilder;
use servers::grpc::flight::FlightCraftRef;
@@ -74,6 +72,7 @@ where
{
pub fn new(opts: T, instance: Arc<Instance>, 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,24 +134,12 @@ 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(),
@@ -439,15 +427,16 @@ 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
}
@@ -459,10 +448,9 @@ fn effective_http_options(opts: &FrontendOptions) -> HttpOptions {
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)
@@ -470,7 +458,19 @@ fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) ->
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));
@@ -515,6 +515,81 @@ 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 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.
@@ -532,7 +607,7 @@ mod tests {
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 +644,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 +654,10 @@ 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::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 +670,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 +772,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),
src/frontend/src/service_config.rs
+1 -1
View File
@@ -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;
src/frontend/src/service_config/pending_rows_batcher.rs
+190 -22
View File
@@ -20,10 +20,28 @@ use serde::{Deserialize, Serialize};
use servers::http::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 {
/// 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<BatcherOptions>,
}
/// Write batching controls shared by HTTP ingestion protocols.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
pub struct BatcherOptions {
/// HTTP write protocols sharing this batcher; empty disables all entrances.
pub protocols: Vec<BatchingProtocol>,
/// Time from the first pending submission to a timed flush. Zero disables batching.
@@ -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,229 @@ 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<Option<BatcherOptions>, D::Error>
where
D: serde::Deserializer<'de>,
{
let options = Option::<BatcherOptions>::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::<FrontendOptions>(&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",
] {
assert!(
toml::from_str::<FrontendOptions>(&format!(
"[pending_rows_batcher.logical_table]\nprotocols = ['{protocol}']"
))
.is_err()
);
}
}
#[test]
fn test_protocols() {
let options: PendingRowsBatcherOptions = toml::from_str(
let options: BatcherOptions = toml::from_str(
"protocols = ['influxdb', 'opentsdb', 'otlp', 'logs', 'loki', 'splunk', 'elasticsearch', 'http_sql', 'prom']",
).unwrap();
assert_eq!(options.protocols.len(), 9);
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::<PendingRowsBatcherOptions>(&format!("protocols = ['{invalid}']"))
.is_err()
toml::from_str::<BatcherOptions>(&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::<PendingRowsBatcherOptions>("flow_notification_queue_capacity = 0")
.is_err()
);
assert!(toml::from_str::<BatcherOptions>("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::<PendingRowsBatcherOptions>(&serialized).unwrap()
toml::from_str::<BatcherOptions>(&serialized).unwrap()
);
}
#[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 +308,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();
src/frontend/src/service_config/prom_store.rs
+1 -1
View File
@@ -54,7 +54,7 @@ fn default_max_concurrent_flushes() -> usize {
}
fn default_worker_channel_capacity() -> usize {
65526
65_536
}
fn default_max_inflight_requests() -> usize {
src/meta-srv/src/metasrv/builder.rs
+1
View File
@@ -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));
src/metric-engine/src/data_region.rs
+1
View File
@@ -227,6 +227,7 @@ impl DataRegion {
| AlterKind::UnsetRegionOptions { keys: _ }
| AlterKind::SetIndexes { options: _ }
| AlterKind::UnsetIndexes { options: _ }
| AlterKind::SetJsonSettings { .. }
| AlterKind::SyncColumns {
column_metadatas: _,
} => {
src/mito-codec/src/row_converter.rs
+30
View File
@@ -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<String>, 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<usize> {
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)),
}
}
}
src/mito-codec/src/row_converter/dense.rs
+378 -103
View File
@@ -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<usize> {
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<usize> {
let invalid = || {
error::InvalidDensePrimaryKeySnafu {
reason: "truncated field or invalid encoding",
fn encoded_len(data_type: &ConcreteDataType, bytes: &[u8]) -> Result<usize> {
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<i64>
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<B: Buf>(
data_type: &ConcreteDataType,
deserializer: &mut Deserializer<B>,
) -> Option<Result<ValueRef<'static>>> {
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<usize> {
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<Item = Result<(ColumnId, Value)>> + '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<u8>,
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<usize>,
deserializer: &mut Deserializer<&[u8]>,
) -> Result<usize> {
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<usize>,
) -> Result<Value> {
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::<Result<Vec<_>>>()
.unwrap();
assert_eq!(
sequential
.iter()
.map(|(_, value)| value)
.collect::<Vec<_>>(),
row.iter().collect::<Vec<_>>()
);
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::<Result<Vec<_>>>()
.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::<Result<Vec<_>>>()
.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::<Result<Vec<_>>>()
.is_err()
);
assert!(
codec
.decode_value_at(&invalid_utf8, 0, &mut Vec::new())
.is_err()
);
}
#[test]
fn test_memcmp_dictionary() {
// Test Dictionary<i32, string>
src/mito-codec/src/row_converter/sparse/checked.rs
+7 -38
View File
@@ -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,
@@ -141,45 +142,13 @@ impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> {
/// Finds the end of an Option<String> 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)]
src/mito2/benches/bench_dense_index_update.rs
+112 -50
View File
@@ -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<Batch> = if encoded_only {
let pk = batch
.column_by_name("__primary_key")
.unwrap()
.as_any()
.downcast_ref::<DictionaryArray<UInt32Type>>()
.unwrap();
let values = pk.values().as_any().downcast_ref::<BinaryArray>().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();
src/mito2/benches/bench_pk_tag_column.rs
+114 -3
View File
@@ -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::<UInt32Type>::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);
src/mito2/src/access_layer.rs
+7 -1
View File
@@ -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<SstInfoArray> {
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();
@@ -426,6 +428,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 +451,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),
src/mito2/src/cache.rs
+142 -68
View File
@@ -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;
@@ -238,7 +240,7 @@ impl SstMetaPreparation {
}
impl CompactSstMeta {
async fn decode(&self) -> Result<Arc<CachedSstMeta>> {
async fn decode(&self, runtime: &Runtime) -> Result<Arc<CachedSstMeta>> {
let mut decoded_guard = self.decoded.lock().await;
if let Some(decoded) = decoded_guard.upgrade() {
return Ok(decoded);
@@ -248,33 +250,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 +288,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<RegionMetadataRef>,
page_index_policy: PageIndexPolicy,
runtime: &Runtime,
) -> Result<Arc<CachedSstMeta>> {
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<RegionMetadataRef>,
page_index_policy: PageIndexPolicy,
runtime: &Runtime,
) -> Result<SstMetaPreparation> {
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 +678,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 +708,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 => {
@@ -1057,6 +1079,7 @@ impl CacheManager {
file_id: RegionFileId,
metrics: &mut MetadataCacheMetrics,
page_index_policy: PageIndexPolicy,
runtime: &Runtime,
) -> Option<Arc<CachedSstMeta>> {
let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id());
let compact = self
@@ -1075,7 +1098,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 +1124,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 +1145,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 +1232,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());
@@ -2144,7 +2167,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 +2233,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 +2268,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 +2290,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(&region_metadata, &cached.region_metadata()));
@@ -2285,9 +2333,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 +2350,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 +2436,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 +2461,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 +2475,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()
);
src/mito2/src/cache/file_cache.rs
Vendored
+5 -2
View File
@@ -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<SstMetaPreparation> {
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<Arc<CachedSstMeta>> {
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
src/mito2/src/compaction/json2.rs
+10 -3
View File
@@ -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),
src/mito2/src/read.rs
+125 -2
View File
@@ -65,7 +65,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};
@@ -119,6 +119,8 @@ pub struct Batch {
primary_key: Vec<u8>,
/// Possibly decoded `primary_key` values. Some places would decode it in advance.
pk_values: Option<CompositeValues>,
/// Lazily decoded Dense columns, shared by consumers such as index creators.
dense_pk_cache: Option<Box<DensePkCache>>,
/// Timestamps of rows, should be sorted and not null.
timestamps: VectorRef,
/// Sequences of rows
@@ -135,6 +137,13 @@ pub struct Batch {
fields_idx: Option<HashMap<ColumnId, usize>>,
}
/// Only used when no fully decoded primary key was supplied by the reader.
#[derive(Debug, PartialEq, Clone)]
struct DensePkCache {
offsets: Vec<usize>,
values: Vec<Option<Value>>,
}
impl Batch {
/// Creates a new batch.
pub fn new(
@@ -173,12 +182,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 +225,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 +281,7 @@ impl Batch {
/// Be sure to update that field as well.
pub fn set_primary_key(&mut self, primary_key: Vec<u8>) {
self.primary_key = primary_key;
self.dense_pk_cache = None;
}
/// Slice the batch, returning a new batch.
@@ -290,6 +303,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)),
@@ -772,9 +786,26 @@ impl Batch {
))
}
/// 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 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 +813,24 @@ impl Batch {
column_id: ColumnId,
) -> Result<Option<&Value>> {
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)?);
}
@@ -1098,6 +1147,7 @@ impl BatchBuilder {
Ok(Batch {
primary_key: self.primary_key,
pk_values: None,
dense_pk_cache: None,
timestamps,
sequences,
op_types,
@@ -1254,6 +1304,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],
src/mito2/src/read/pruner.rs
+7 -5
View File
@@ -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 {
src/mito2/src/region/opener.rs
+7 -1
View File
@@ -1247,7 +1247,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 +1302,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
{
src/mito2/src/sst/index.rs
+70 -3
View File
@@ -37,11 +37,13 @@ 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;
@@ -57,8 +59,8 @@ use crate::cache::{CacheManagerRef, CacheStrategy};
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,
@@ -272,6 +274,8 @@ pub struct Indexer {
last_mem_fulltext_index: usize,
bloom_filter_indexer: Option<BloomFilterIndexer>,
last_mem_bloom_filter: usize,
/// Present only when the active creators together need every Dense PK field.
dense_pk_decoder: Option<DensePrimaryKeyCodec>,
#[cfg(feature = "vector_index")]
vector_indexer: Option<VectorIndexer>,
#[cfg(feature = "vector_index")]
@@ -280,6 +284,54 @@ pub struct Indexer {
}
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<InvertedIndexer>,
bloom_filter_indexer: Option<BloomFilterIndexer>,
) -> 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;
@@ -398,6 +450,7 @@ 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());
indexer.prepare_dense_pk_decoder(&self.metadata);
#[cfg(feature = "vector_index")]
{
indexer.vector_indexer = self.build_vector_indexer(region_file_id.file_id());
@@ -774,7 +827,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!(
src/mito2/src/sst/index/column_test.rs
+85 -5
View File
@@ -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;
@@ -201,13 +203,85 @@ async fn materialized_runs_match_legacy_index_bytes_and_bloom_lookups() {
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 +289,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" => {
src/mito2/src/sst/index/indexer/abort.rs
+1
View File
@@ -20,6 +20,7 @@ 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;
src/mito2/src/sst/index/indexer/finish.rs
+1
View File
@@ -27,6 +27,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 {
src/mito2/src/sst/index/indexer/update.rs
+22
View File
@@ -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;
}
@@ -39,6 +44,23 @@ impl Indexer {
}
}
/// 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.
async fn do_update_inverted_index(&mut self, batch: &mut Batch) -> bool {
let Some(creator) = self.inverted_indexer.as_mut() else {
src/mito2/src/sst/parquet/flat_format.rs
+220 -27
View File
@@ -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};
@@ -676,8 +676,9 @@ pub(crate) fn sst_column_id_indices(metadata: &RegionMetadata) -> HashMap<Column
/// Decodes primary keys from a batch and returns decoded primary key information.
///
/// The batch must contain a primary key column at the expected index.
/// Sparse primary keys stay encoded: tag values are extracted lazily per column
/// Primary keys stay encoded: tag values are extracted lazily per column
/// in [`DecodedPrimaryKeys::get_tag_column`] without decoding every label.
/// The codec must describe the source key's field order and types.
pub fn decode_primary_keys(
codec: &dyn PrimaryKeyCodec,
batch: &RecordBatch,
@@ -729,15 +730,17 @@ pub fn decode_primary_keys(
values: pk_values_array.clone(),
distinct_keys,
},
PrimaryKeyEncoding::Dense => {
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 +751,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<CompositeValues>),
/// Dense keys share positional offsets across projected tag columns.
Dense {
codec: DensePrimaryKeyCodec,
values: BinaryArray,
distinct_keys: Vec<u32>,
/// Initialized by positional access; full extraction needs no offsets.
offsets: Vec<Vec<usize>>,
},
/// Sparse primary keys stay encoded; each tag column extracts only its own
/// values from the raw keys.
Sparse {
@@ -811,7 +820,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 +836,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 +888,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<Vec<ArrayRef>> {
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 +1063,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 +1147,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::<UInt32Type>::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::<UInt32Type>::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,
src/mito2/src/sst/parquet/json_align.rs
+75 -3
View File
@@ -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<ArrayRef> {
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::<Int64Array>()
.unwrap()
.value(0)
);
assert_eq!(
20,
output
.column(1)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.value(1)
);
}
#[test]
fn test_rewrite_rejects_mismatched_output_layout() {
let columns = HashMap::from([(
src/mito2/src/sst/parquet/reader.rs
+8 -1
View File
@@ -851,7 +851,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) => {
src/mito2/src/test_util/bench_util.rs
+19 -1
View File
@@ -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<ColumnId>,
file_schema: SchemaRef,
) -> impl Fn(RecordBatch) -> crate::error::Result<RecordBatch> {
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,
src/object-store/src/test_util.rs
-29
View File
@@ -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<String>,
}
/// Returns s3 test config, return None if not found.
pub fn s3_test_config() -> Option<TestS3Config> {
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
}
src/object-store/tests/object_store_test.rs
+9 -1
View File
@@ -284,13 +284,21 @@ 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);
// 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);
src/operator/src/batcher.rs
+2 -1
View File
@@ -29,7 +29,8 @@ pub trait PendingRowsBatcher: Send + Sync {
/// Acquires one slot per original request, shared by all of its table submissions.
async fn acquire(&self) -> Result<Arc<OwnedSemaphorePermit>>;
/// 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.
async fn submit(
&self,
src/operator/src/expr_helper.rs
+84 -25
View File
@@ -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,
@@ -782,6 +784,29 @@ pub(crate) fn to_repartition_request(
})
}
fn json_settings_to_proto(settings: JsonSettings) -> Result<PbJsonSettings> {
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::<Result<Vec<_>>>()?;
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,
@@ -830,12 +855,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 +881,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(),
@@ -1752,31 +1794,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]
src/operator/src/insert.rs
+310 -5
View File
@@ -156,6 +156,127 @@ 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<bool> {
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.
// 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::TimestampMillisecond
| 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<u64> {
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,
@@ -2313,19 +2434,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::<Vec<_>>();
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::<Vec<_>>();
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]
@@ -2570,6 +2739,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<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
src/query/src/promql/planner.rs
+36 -7
View File
@@ -13515,6 +13515,30 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom
}
}
#[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.
@@ -13539,13 +13563,18 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom
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}"
);
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]
src/query/src/promql/planner/matching_filters.rs
+202 -98
View File
@@ -52,26 +52,23 @@ 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, so it does not affect a one-to-one cardinality check such as #9209.
/// Grouped matching is conservatively restricted to a provably unique one-side operand until
/// its cardinality checks are implemented (#9208).
///
/// `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 +105,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 +128,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 +164,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 +188,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 +204,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<bool> {
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 +297,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<Expr>]) -> Option<usize> {
let mut found = None;
for (index, arg) in args.iter().enumerate() {
@@ -322,7 +350,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 +518,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])"#,
src/servers/src/batcher.rs
+18
View File
@@ -21,3 +21,21 @@ pub mod table;
#[cfg(test)]
mod test_util;
/// 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::<bool>().ok())
.unwrap_or(true)
}
src/servers/src/batcher/logical_table.rs
+23 -22
View File
@@ -21,6 +21,7 @@ mod tables;
#[cfg(test)]
mod test_util;
use std::future::{Future, ready};
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Duration;
@@ -54,6 +55,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,24 +64,6 @@ 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::<bool>().ok())
.unwrap_or(true)
}
const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3;
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
@@ -182,14 +166,31 @@ impl LogicalTablePendingRowsBatcher {
impl LogicalTablePendingRowsBatcher {
pub async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result<u64> {
self.submit_with(requests, ctx, |_| ready(Ok(())))
.await
.map(|(rows, ())| rows)
}
/// Submits with request-level accounting after schema preparation and before
/// queue admission. Acknowledgement follows the global batching policy.
pub async fn submit_with<T, F>(
&self,
requests: RowInsertRequests,
ctx: QueryContextRef,
after_prepare: impl FnOnce(RowInsertRequests) -> F,
) -> Result<(u64, T)>
where
F: Future<Output = Result<T>> + 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 +270,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))
}
}
}
src/servers/src/batcher/logical_table/tables.rs
+22 -22
View File
@@ -75,7 +75,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();
@@ -113,13 +113,13 @@ 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 +127,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 +137,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<Vec<(&str, &[ColumnSchema])>> {
pub(in crate::batcher::logical_table) fn collect_unique_table_schemas<'a>(
table_rows: &[(&'a str, &'a Rows)],
) -> Result<Vec<(&'a str, &'a [ColumnSchema])>> {
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 +233,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 +298,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 +308,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 +321,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 +397,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<String, (Arc<ArrowSchema>, u32)>,
) -> Result<Vec<(String, u32, RecordBatchWithTsIdx)>> {
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 +415,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 +531,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());
src/servers/src/batcher/table.rs
+114 -3
View File
@@ -38,6 +38,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 +71,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<Inserter>,
flow_notifier: FlowNotifier,
@@ -110,6 +112,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,8 +163,8 @@ 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,
@@ -210,6 +213,9 @@ impl PendingRowsBatcher for TablePendingRowsBatcher {
}
.fail();
}
if !self.pending_rows_batch_sync {
return Ok(total_rows);
}
response_rx
.await
.map_err(|_| {
@@ -236,7 +242,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 +326,111 @@ 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);
if sync {
assert!(!submitted.as_ref().unwrap().is_finished());
} else {
assert_eq!(
timeout(Duration::from_secs(5), submitted.take().unwrap())
.await
.unwrap()
.unwrap()
.unwrap(),
1
);
}
// 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 {
assert_eq!(result.unwrap(), 1);
}
}
timeout(Duration::from_secs(5), acquire)
.await
.unwrap()
.unwrap();
}
}
}
fn rows(value: i32) -> Rows {
Rows {
schema: vec![
src/servers/src/batcher/table/batch.rs
+2
View File
@@ -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)));
src/servers/src/http.rs
+62 -1
View File
@@ -215,6 +215,7 @@ pub struct HttpServer {
// server configs
options: HttpOptions,
batching_protocols: Vec<BatchingProtocol>,
logical_batching_protocols: Vec<BatchingProtocol>,
bind_addr: Option<SocketAddr>,
/// What this server instance exposes. See [`HttpServerKind`].
kind: HttpServerKind,
@@ -235,6 +236,9 @@ 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<BatchingProtocol>);
/// Sets a local-only write selector after authentication creates the context.
async fn set_http_write_batching(
State(protocol): State<BatchingProtocol>,
@@ -245,8 +249,13 @@ async fn set_http_write_batching(
.extensions()
.get::<Arc<Vec<BatchingProtocol>>>()
.is_some_and(|protocols| protocols.contains(&protocol));
let logical_enabled = req
.extensions()
.get::<LogicalBatchingProtocols>()
.is_some_and(|protocols| protocols.0.contains(&protocol));
if let Some(ctx) = req.extensions_mut().get_mut::<QueryContext>() {
ctx.set_batching_enabled(enabled);
ctx.set_logical_batching_enabled(logical_enabled);
}
next.run(req).await
}
@@ -650,6 +659,7 @@ pub struct DashboardState {
pub struct HttpServerBuilder {
options: HttpOptions,
batching_protocols: Vec<BatchingProtocol>,
logical_batching_protocols: Vec<BatchingProtocol>,
user_provider: Option<UserProviderRef>,
router: Router,
memory_limiter: ServerMemoryLimiter,
@@ -660,13 +670,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<BatchingProtocol>) -> 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<BatchingProtocol>) -> Self {
self.batching_protocols = protocols;
self
@@ -894,6 +911,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 +946,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 +965,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),
@@ -1104,6 +1124,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(
@@ -2462,6 +2485,44 @@ mod batching_tests {
use crate::opentsdb::codec::DataPoint;
use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler};
#[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<QueryContext>| 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]);
}
}
}
#[test]
fn test_protocol_names_reject_unknown_values() {
assert_eq!(
src/servers/src/postgres/types.rs
+308 -7
View File
@@ -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<S: Encoder>(
_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<S: Encoder>(
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<S, B>
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::<i32>();
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<Type> {
&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 {
@@ -1921,6 +1946,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<DataRow> = 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<Option<String>> = 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<DataRow> = 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<Option<String>> = 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 = [
src/session/src/context.rs
+22 -2
View File
@@ -82,6 +82,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,
@@ -458,7 +461,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 HTTP entry point selected ordinary-table batching.
pub fn batching_enabled(&self) -> bool {
self.batching_enabled
}
@@ -634,6 +647,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(),
@@ -923,11 +937,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::from(wire);
assert!(!restored.batching_enabled());
assert!(!restored.logical_batching_enabled());
let ctx = QueryContextBuilder::default()
.set_extension("batching_enabled".to_string(), "true".to_string())
.build();
src/sql/src/parsers/alter_parser.rs
+13 -15
View File
@@ -462,19 +462,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,
})
}
}
@@ -1576,9 +1576,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 +1585,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 +1608,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!()
};
src/sql/src/statements/alter.rs
+37 -1
View File
@@ -88,6 +88,11 @@ pub enum AlterTableOperation {
target_type: DataType,
json2_options: Option<Json2Options>,
},
/// `MODIFY <column_name> JSON2 [json2_options]`
SetJsonSettings {
column_name: Ident,
json2_options: Option<Json2Options>,
},
/// `SET <table attrs key> = <table attr value>`
SetTableOptions {
options: Vec<KeyValueOption>,
@@ -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())
src/sql/src/statements/create.rs
+21 -18
View File
@@ -157,6 +157,26 @@ impl Display for Json2Options {
}
}
impl Json2Options {
pub fn build_json_settings(&self) -> Result<JsonSettings> {
let type_hints = self
.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::<Result<Vec<_>>>()?;
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)
}
}
#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)]
pub struct JsonTypeHint {
pub path: Vec<String>,
@@ -351,24 +371,7 @@ impl ColumnExtensions {
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::<Result<Vec<_>>>()?;
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<()> {
src/standalone/src/options.rs
+68 -5
View File
@@ -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",
])
}
}
@@ -219,6 +220,61 @@ mod tests {
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::<StandaloneOptions>(
"[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![
servers::http::BatchingProtocol::Prom,
servers::http::BatchingProtocol::Otlp
]
);
},
);
}
#[test]
fn test_batcher_protocols_from_env() {
temp_env::with_vars(
@@ -231,7 +287,7 @@ mod tests {
StandaloneOptions::load_layered_options(None, "STANDALONE_BATCHER_TEST")
.unwrap();
assert_eq!(
options.pending_rows_batcher.protocols,
options.pending_rows_batcher.table.protocols,
vec![
servers::http::BatchingProtocol::Influxdb,
servers::http::BatchingProtocol::HttpSql
@@ -247,6 +303,7 @@ mod tests {
assert!(
!defaults
.pending_rows_batcher
.table
.pending_rows_batching_enabled()
);
let options: StandaloneOptions = toml::from_str(
@@ -259,16 +316,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);
src/store-api/Cargo.toml
+1
View File
@@ -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
src/store-api/src/metadata.rs
+36
View File
@@ -30,6 +30,7 @@ use common_error::status_code::StatusCode;
use common_macro::stack_trace_debug;
use datatypes::arrow;
use datatypes::arrow::datatypes::FieldRef;
use datatypes::extension::json::json2_metadata_with_updated_settings;
use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef, VectorIndexOptions};
use datatypes::types::TimestampType;
use itertools::Itertools;
@@ -663,6 +664,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 +836,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<SetIndexOption>) -> Result<()> {
let mut set_index_map: HashMap<_, Vec<_>> = HashMap::new();
for option in &options {
src/store-api/src/region_request.rs
+166
View File
@@ -32,6 +32,7 @@ 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 +40,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;
@@ -795,6 +798,13 @@ pub enum AlterKind {
/// Columns to change.
columns: Vec<ModifyColumnType>,
},
/// Set JSON2 settings of a region column.
SetJsonSettings {
/// Column name.
column_name: String,
/// Target JSON2 settings.
settings: JsonSettings,
},
/// Set region options.
SetRegionOptions { options: Vec<SetRegionOption> },
/// Unset region options.
@@ -975,6 +985,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 +1097,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::<Json2ExtensionType>()
.ok()
.flatten()
.is_none_or(|extension| {
!extension.metadata().json_settings().equivalent(settings)
})
}),
AlterKind::SetRegionOptions { .. } => true,
AlterKind::UnsetRegionOptions { .. } => true,
AlterKind::SetIndexes { options, .. } => options
@@ -1160,6 +1185,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<alter_request::Kind> for AlterKind {
@@ -1183,6 +1236,15 @@ impl TryFrom<alter_request::Kind> for AlterKind {
.collect::<Vec<_>>();
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 +1546,36 @@ impl From<v1::ModifyColumnType> for ModifyColumnType {
}
}
fn json_settings_from_proto(settings: v1::JsonSettings) -> Result<JsonSettings> {
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::<Result<Vec<_>>>()?;
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
@@ -1953,8 +2045,11 @@ 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;
@@ -2364,6 +2459,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 +2569,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 +2862,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 {
src/table/src/metadata.rs
+163 -3
View File
@@ -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),
@@ -1187,6 +1189,77 @@ impl TableMeta {
Ok(meta_builder)
}
fn set_json_settings(
&self,
table_name: &str,
req: &SetJsonSettingsRequest,
) -> Result<TableMetaBuilder> {
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());
src/table/src/requests.rs
+11
View File
@@ -25,6 +25,7 @@ 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,
@@ -371,6 +372,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 +669,9 @@ pub enum AlterKind {
ModifyColumnTypes {
columns: Vec<ModifyColumnTypeRequest>,
},
SetJsonSettings {
request: SetJsonSettingsRequest,
},
RenameTable {
new_table_name: String,
},
tests-fuzz/README.md
+6
View File
@@ -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.
tests-integration/Cargo.toml
+1 -2
View File
@@ -65,7 +65,7 @@ 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
@@ -111,7 +111,6 @@ hex.workspace = true
http.workspace = true
itertools.workspace = true
jsonb.workspace = true
meta-srv = { workspace = true, features = ["mock"] }
mito2 = { workspace = true, features = ["test", "test-shared-fs-region-migration"] }
mysql_async = { version = "0.37", default-features = false, features = [
"time",
tests-integration/src/cluster.rs
+14
View File
@@ -36,6 +36,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};
@@ -174,6 +175,7 @@ pub struct GreptimeDbClusterBuilder {
shared_home_dir: Option<Arc<TempDir>>,
meta_selector: Option<SelectorRef>,
local_file_access: LocalFileAccess,
event_recorder_options: EventRecorderOptions,
}
impl GreptimeDbClusterBuilder {
@@ -209,6 +211,7 @@ impl GreptimeDbClusterBuilder {
shared_home_dir: None,
meta_selector: None,
local_file_access: LocalFileAccess::default(),
event_recorder_options: EventRecorderOptions::default(),
}
}
@@ -253,6 +256,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 {
@@ -303,6 +315,7 @@ impl GreptimeDbClusterBuilder {
..Default::default()
},
gc: self.metasrv_gc_config.clone(),
event_recorder: self.event_recorder_options.clone(),
..Default::default()
};
@@ -550,6 +563,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()
};
tests-integration/src/otlp.rs
+227
View File
@@ -503,6 +503,233 @@ WITH(
Ok(())
}
#[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::pending_rows_batch_sync_enabled;
use servers::http::BatchingProtocol;
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 mut options = standalone.opts.clone();
options.otlp.experimental_enable_exponential_histogram = true;
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::<usize>() == 2
{
break batches;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap();
assert_eq!(
batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
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::<usize>(),
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")
tests-integration/src/standalone.rs
+14
View File
@@ -52,6 +52,7 @@ 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;
@@ -88,6 +89,7 @@ pub struct GreptimeDbStandaloneBuilder {
event_recorder_options: EventRecorderOptions,
auto_create_table: bool,
experimental_metric_export: bool,
logical_batcher: Option<BatcherOptions>,
}
impl GreptimeDbStandaloneBuilder {
@@ -108,6 +110,7 @@ impl GreptimeDbStandaloneBuilder {
event_recorder_options: EventRecorderOptions::default(),
auto_create_table: true,
experimental_metric_export: false,
logical_batcher: None,
}
}
@@ -118,6 +121,13 @@ impl GreptimeDbStandaloneBuilder {
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;
@@ -377,6 +387,10 @@ impl GreptimeDbStandaloneBuilder {
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(),
..Default::default()
},
// Tests cover the descriptor, so they run with it enabled.
otlp: frontend::service_config::OtlpOptions {
experimental_enable_resource_info: true,
tests-integration/src/test_util.rs
+110 -24
View File
@@ -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_path>/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<TempFolder>,
/// Local read cache directory. Only held so it is removed when the guard drops.
#[allow(dead_code)]
local_cache: Option<TempDir>,
}
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);
}
@@ -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<Box<dyn Server>>) {
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, true)
.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::<SocketAddr>().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()

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