diff --git a/.github/actions/build-linux-artifacts/action.yml b/.github/actions/build-linux-artifacts/action.yml
index 3cb9c439551..c5212c28998 100644
--- a/.github/actions/build-linux-artifacts/action.yml
+++ b/.github/actions/build-linux-artifacts/action.yml
@@ -31,6 +31,22 @@ inputs:
description: Build GreptimeDB with large page size (65536).
required: false
default: 'false'
+ aws-ci-test-bucket:
+ description: Bucket used by the object-store integration tests
+ required: false
+ default: ''
+ aws-region:
+ description: Region of the test bucket
+ required: false
+ default: ''
+ aws-access-key-id:
+ description: Access key id for the test bucket
+ required: false
+ default: ''
+ aws-secret-access-key:
+ description: Secret access key for the test bucket
+ required: false
+ default: ''
runs:
using: composite
steps:
@@ -38,6 +54,13 @@ runs:
if: ${{ inputs.disable-run-tests == 'false' }}
shell: bash
# NOTE: If the BUILD_JOBS > 4, it's always OOM in EC2 instance.
+ # Without GT_S3_*, every remote object-store case is skipped and the
+ # release only ever exercises the local file backend.
+ env:
+ GT_S3_BUCKET: ${{ inputs.aws-ci-test-bucket }}
+ GT_S3_REGION: ${{ inputs.aws-region }}
+ GT_S3_ACCESS_KEY_ID: ${{ inputs.aws-access-key-id }}
+ GT_S3_ACCESS_KEY: ${{ inputs.aws-secret-access-key }}
run: |
cd ${{ inputs.working-dir }} && \
make run-it-in-container BUILD_JOBS=4 \
diff --git a/.github/scripts/run-fuzz-targets-test.sh b/.github/scripts/run-fuzz-targets-test.sh
index 5e4cda64366..9705b6bc5af 100755
--- a/.github/scripts/run-fuzz-targets-test.sh
+++ b/.github/scripts/run-fuzz-targets-test.sh
@@ -31,7 +31,7 @@ new_fixture() {
#!/usr/bin/env bash
set -euo pipefail
printf '%s\t%s\t%s\n' "${GT_FUZZ_DUMP_DIR}" "$*" "${MOCK_CARGO_MARKER:-}" >>"${MOCK_CARGO_LOG}"
-target="$3"
+target="$4"
case " ${MOCK_FAIL_TARGETS:-} " in
*" ${target} "*) exit 17 ;;
esac
@@ -117,6 +117,20 @@ EOF
set -e
}
+test_rustup_toolchain_env_is_honored() {
+ new_fixture
+ export RUSTUP_TOOLCHAIN=nightly-2026-03-21
+ run_fixture $'fuzz_create_table' true true "fuzz_create_table"
+ unset RUSTUP_TOOLCHAIN
+
+ assert_eq 1 "${fixture_status}" "pinned toolchain run status"
+ grep -q -- '+nightly-2026-03-21 fuzz run fuzz_create_table' "${fixture}/cargo.log" || \
+ fail "pinned toolchain missing from cargo invocation"
+ grep -q 'cargo +nightly-2026-03-21 fuzz run fuzz_create_table' "${fixture}/artifacts/summary.md" || \
+ fail "pinned toolchain missing from reproduce command"
+ cleanup_fixture
+}
+
test_successful_targets_run_in_order() {
new_fixture
run_fixture $'fuzz_create_table\nfuzz_insert' false true ""
@@ -125,7 +139,7 @@ test_successful_targets_run_in_order() {
assert_eq 2 "$(wc -l <"${fixture}/cargo.log" | tr -d ' ')" "cargo invocation count"
assert_eq \
$'fuzz_create_table\nfuzz_insert' \
- "$(awk -F '\t' '{print $2}' "${fixture}/cargo.log" | sed -E 's/^fuzz run ([^ ]+).*/\1/')" \
+ "$(awk -F '\t' '{print $2}' "${fixture}/cargo.log" | sed -E 's/^\+nightly fuzz run ([^ ]+).*/\1/')" \
"target order"
grep -q -- '--features=unstable' "${fixture}/cargo.log" || fail "unstable feature missing"
grep -q -- '-max_total_time=120' "${fixture}/cargo.log" || fail "fuzz time missing"
@@ -175,7 +189,7 @@ test_fail_fast_stops_after_first_failure() {
fail "skipped target annotation missing"
grep -q '### Reproduce failed targets' "${fixture}/artifacts/summary.md" || \
fail "reproduction section missing"
- grep -q 'cargo fuzz run fuzz_insert' "${fixture}/artifacts/summary.md" || \
+ grep -q 'cargo +nightly fuzz run fuzz_insert' "${fixture}/artifacts/summary.md" || \
fail "reproduction command missing"
cleanup_fixture
}
@@ -367,5 +381,6 @@ test_collector_keeps_target_scopes_separate
test_cluster_collector_honors_target_scope_and_namespace
test_setup_failure_writes_artifact_contract
test_setup_failure_keeps_manifest_when_collection_fails
+test_rustup_toolchain_env_is_honored
printf 'All fuzz orchestration script tests passed.\n'
diff --git a/.github/scripts/run-fuzz-targets.sh b/.github/scripts/run-fuzz-targets.sh
index 226a393e16c..205ba58c419 100755
--- a/.github/scripts/run-fuzz-targets.sh
+++ b/.github/scripts/run-fuzz-targets.sh
@@ -104,7 +104,7 @@ write_summary() {
reproduce_args+=("GT_FUZZ_INSTANCE_ROOT_DIR=${GT_FUZZ_INSTANCE_ROOT_DIR}")
fi
fi
- reproduce_args+=(cargo fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none)
+ reproduce_args+=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none)
if [[ "${fuzz_unstable}" == true ]]; then
reproduce_args+=(--features=unstable)
fi
@@ -195,7 +195,7 @@ for ((index = 0; index < ${#targets[@]}; index++)); do
fi
run_args=("${fuzz_binary}" "-max_total_time=${FUZZ_MAX_TOTAL_TIME}" "-artifact_prefix=${target_dir}/libfuzzer/")
else
- run_args=(cargo fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none)
+ run_args=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none)
if [[ "${fuzz_unstable}" == true ]]; then
run_args+=(--features=unstable)
fi
diff --git a/.github/workflows/dev-build.yml b/.github/workflows/dev-build.yml
index d03fbeff14a..563435b708e 100644
--- a/.github/workflows/dev-build.yml
+++ b/.github/workflows/dev-build.yml
@@ -183,6 +183,10 @@ jobs:
image-registry: ${{ vars.ECR_IMAGE_REGISTRY }}
image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }}
large-page-size: ${{ inputs.large-page-size }}
+ aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }}
+ aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }}
+ aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }}
+ aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }}
build-linux-arm64-artifacts:
name: Build linux-arm64 artifacts
@@ -217,6 +221,10 @@ jobs:
image-registry: ${{ vars.ECR_IMAGE_REGISTRY }}
image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }}
large-page-size: ${{ inputs.large-page-size }}
+ aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }}
+ aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }}
+ aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }}
+ aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }}
release-images-to-dockerhub:
name: Build and push images to DockerHub
diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml
index c81dba199e9..4e0634fbc57 100644
--- a/.github/workflows/integration.yml
+++ b/.github/workflows/integration.yml
@@ -31,6 +31,11 @@ name: Integration CI
env:
CARGO_FUZZ_VERSION: "0.13.2"
+ # Fuzz targets require a nightly toolchain (cargo-fuzz needs
+ # `-Zsanitizer=fuzzer`). This pin decouples the fuzz workflow from the
+ # workspace default in rust-toolchain.toml, so it keeps working when the
+ # workspace moves to a stable toolchain.
+ FUZZ_RUST_TOOLCHAIN: "nightly-2026-03-21"
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -328,6 +333,7 @@ jobs:
- uses: actions-rust-lang/setup-rust-toolchain@v1
with:
cache: false
+ toolchain: ${{ env.FUZZ_RUST_TOOLCHAIN }}
- name: Fuzz Dependency Cache
id: fuzz-dependencies-cache
uses: Swatinem/rust-cache@v2
@@ -343,7 +349,7 @@ jobs:
uses: actions/cache/restore@v4
with:
path: ${{ runner.temp }}/greptime-fuzz-binaries
- key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', 'rust-toolchain.toml', '.cargo/**', 'tests-fuzz/**', 'src/**') }}
+ key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }}
- name: Report fuzz cache lookup
shell: bash
run: |
@@ -355,6 +361,8 @@ jobs:
- name: Set Rust Fuzz
if: steps.fuzz-binaries.outputs.cache-hit != 'true'
shell: bash
+ env:
+ RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }}
run: |
sudo apt-get install -y libfuzzer-14-dev
cargo install cargo-fuzz --version "${CARGO_FUZZ_VERSION}" --locked
@@ -363,6 +371,7 @@ jobs:
shell: bash
env:
CUSTOM_LIBFUZZER_PATH: /usr/lib/llvm-14/lib/libFuzzer.a
+ RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }}
FUZZ_TARGETS: |
fuzz_create_database
fuzz_create_table
@@ -410,7 +419,7 @@ jobs:
uses: actions/cache/save@v4
with:
path: ${{ runner.temp }}/greptime-fuzz-binaries
- key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', 'rust-toolchain.toml', '.cargo/**', 'tests-fuzz/**', 'src/**') }}
+ key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }}
- name: Pack prebuilt fuzz binaries
shell: bash
run: |
@@ -479,6 +488,7 @@ jobs:
- name: Fuzz Test
uses: ./.github/actions/fuzz-test
env:
+ RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }}
GT_MYSQL_ADDR: 127.0.0.1:4002
FUZZ_SERVICE_LOG: /tmp/greptime-fuzz-runtime/greptime.log
with:
@@ -537,6 +547,7 @@ jobs:
- name: Run Fuzz Test
uses: ./.github/actions/fuzz-test
env:
+ RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }}
GT_MYSQL_ADDR: 127.0.0.1:4002
GT_FUZZ_BINARY_PATH: ./bin/greptime
GT_FUZZ_INSTANCE_ROOT_DIR: /tmp/greptime-fuzz-artifacts/targets/unstable_fuzz_create_table_standalone/service/instance/
@@ -820,6 +831,7 @@ jobs:
- name: Fuzz Test
uses: ./.github/actions/fuzz-test
env:
+ RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }}
GT_MYSQL_ADDR: 127.0.0.1:4002
GT_KAFKA_WAL_HELPER_URL: http://127.0.0.1:8080
GT_KAFKA_ADDR: kafka.kafka-cluster.svc.cluster.local:9092
@@ -994,6 +1006,7 @@ jobs:
- name: Fuzz Test
uses: ./.github/actions/fuzz-test
env:
+ RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }}
GT_MYSQL_ADDR: 127.0.0.1:4002
with:
targets: ${{ matrix.group.targets }}
diff --git a/.github/workflows/nightly-build.yml b/.github/workflows/nightly-build.yml
index 14ebb6e715e..c78ca4dc5b9 100644
--- a/.github/workflows/nightly-build.yml
+++ b/.github/workflows/nightly-build.yml
@@ -151,6 +151,10 @@ jobs:
disable-run-tests: ${{ env.DISABLE_RUN_TESTS }}
image-registry: ${{ vars.ECR_IMAGE_REGISTRY }}
image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }}
+ aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }}
+ aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }}
+ aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }}
+ aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }}
build-linux-arm64-artifacts:
name: Build linux-arm64 artifacts
@@ -173,6 +177,10 @@ jobs:
disable-run-tests: ${{ env.DISABLE_RUN_TESTS }}
image-registry: ${{ vars.ECR_IMAGE_REGISTRY }}
image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }}
+ aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }}
+ aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }}
+ aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }}
+ aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }}
release-images-to-dockerhub:
name: Build and push images to DockerHub
diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml
index d4afa2778b8..8441ed02312 100644
--- a/.github/workflows/nightly-ci.yml
+++ b/.github/workflows/nightly-ci.yml
@@ -71,6 +71,57 @@ jobs:
path: C:\Users\RUNNER~1\AppData\Local\Temp\sqlness*
retention-days: 3
+ # PR CI runs the object-store matrix against a local MinIO. This job is where
+ # the same matrix meets real AWS S3.
+ test-on-linux-real-s3:
+ name: Run tests on Linux against real S3
+ if: ${{ github.repository == 'GreptimeTeam/greptimedb' }}
+ runs-on: ubuntu-22.04-arm
+ timeout-minutes: 90
+ steps:
+ - uses: actions/checkout@v4
+ with:
+ persist-credentials: false
+ - uses: arduino/setup-protoc@v3
+ with:
+ repo-token: ${{ secrets.GITHUB_TOKEN }}
+ - uses: rui314/setup-mold@v1
+ - name: Install toolchain
+ uses: actions-rust-lang/setup-rust-toolchain@v1
+ with:
+ cache: false
+ - name: Rust Cache
+ uses: Swatinem/rust-cache@v2
+ with:
+ shared-key: "coverage-test"
+ cache-all-crates: "true"
+ save-if: false
+ - name: Install latest nextest release
+ uses: taiki-e/install-action@nextest
+ - name: Setup external services
+ working-directory: tests-integration/fixtures
+ run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait
+ - name: Run nextest cases
+ run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F 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:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index acf52e241c2..6c0c36c895e 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -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
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 21a6b8d678f..d4b76b8738f 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -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
diff --git a/Cargo.lock b/Cargo.lock
index 2c7f3333329..ea55f1b297d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -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",
diff --git a/Cargo.toml b/Cargo.toml
index b4c4c9d7f54..7c5c53d91c8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -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"
diff --git a/Makefile b/Makefile
index 2526f3e31eb..7176eafd0bb 100644
--- a/Makefile
+++ b/Makefile
@@ -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}
diff --git a/config/config.md b/config/config.md
index 16384b406f4..d736c24d848 100644
--- a/config/config.md
+++ b/config/config.md
@@ -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.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. |
+| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. |
| `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions |
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
@@ -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.
Available values: "last_non_null", "last_row". |
-| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. |
+| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in 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.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. |
| `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. |
| `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. |
-| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. |
+| `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. |
| `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. |
| `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. |
+| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. |
+| `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. |
+| `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. |
+| `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. |
+| `pending_rows_batcher.logical_table.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per physical-table worker. |
+| `pending_rows_batcher.logical_table.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. |
+| `pending_rows_batcher.logical_table.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued logical-table Flow notifications. |
| `jaeger` | -- | -- | Jaeger protocol options. |
| `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. |
| `otlp` | -- | -- | OpenTelemetry protocol options. |
@@ -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.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. |
| `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. |
-| `prom_store.pending_rows_flush_interval` | String | `0s` | Interval to flush pending rows batcher.
Set to "0s" to disable batching mode in Prometheus Remote Write endpoint |
-| `prom_store.max_batch_rows` | Integer | `100000` | Max rows per pending batch before triggering a flush. |
-| `prom_store.max_concurrent_flushes` | Integer | `256` | Max number of concurrent batch flushes. |
-| `prom_store.worker_channel_capacity` | Integer | `65526` | Capacity of the pending batch worker channel. |
-| `prom_store.max_inflight_requests` | Integer | `3000` | Max inflight write requests before backpressure. |
-| `prom_store.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of logical-table flow notifications waiting in the shared queue. |
| `wal` | -- | -- | The WAL options. |
| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** |
| `wal.dir` | String | Unset | The directory to store the WAL files.
**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.
Defaults to max(num_cpus / 2, 2). |
| `http` | -- | -- | The HTTP server options. |
| `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. |
-| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. |
+| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. |
| `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions |
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
@@ -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.
Available values: "last_non_null", "last_row". |
-| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. |
+| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in 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.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. |
| `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. |
| `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. |
-| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. |
+| `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. |
| `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. |
| `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. |
+| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. |
+| `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. |
+| `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. |
+| `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. |
+| `pending_rows_batcher.logical_table.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per physical-table worker. |
+| `pending_rows_batcher.logical_table.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. |
+| `pending_rows_batcher.logical_table.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued logical-table Flow notifications. |
| `jaeger` | -- | -- | Jaeger protocol options. |
| `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. |
| `otlp` | -- | -- | OpenTelemetry protocol options. |
@@ -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.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. |
| `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. |
-| `prom_store.pending_rows_flush_interval` | String | `0s` | Interval to flush pending rows batcher.
Set to "0s" to disable batching mode in Prometheus Remote Write endpoint |
-| `prom_store.max_batch_rows` | Integer | `100000` | Max rows per pending batch before triggering a flush. |
-| `prom_store.max_concurrent_flushes` | Integer | `256` | Max number of concurrent batch flushes. |
-| `prom_store.worker_channel_capacity` | Integer | `65526` | Capacity of the pending batch worker channel. |
-| `prom_store.max_inflight_requests` | Integer | `3000` | Max inflight write requests before backpressure. |
-| `prom_store.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of logical-table flow notifications waiting in the shared queue. |
| `meta_client` | -- | -- | The metasrv client options. |
| `meta_client.metasrv_addrs` | Array | -- | The addresses of the metasrv. |
| `meta_client.timeout` | String | `3s` | Operation timeout. |
diff --git a/config/frontend.example.toml b/config/frontend.example.toml
index 043da57db74..275db3105c5 100644
--- a/config/frontend.example.toml
+++ b/config/frontend.example.toml
@@ -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]
diff --git a/config/standalone.example.toml b/config/standalone.example.toml
index e59e7ba1b4c..257f48eda2c 100644
--- a/config/standalone.example.toml
+++ b/config/standalone.example.toml
@@ -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]
diff --git a/docs/how-to/ai-functions.md b/docs/how-to/ai-functions.md
new file mode 100644
index 00000000000..5c33ffff50e
--- /dev/null
+++ b/docs/how-to/ai-functions.md
@@ -0,0 +1,247 @@
+# AI functions (experimental MVP)
+
+Three asynchronous SQL scalar functions provide natural-language matching,
+classification, and rating. The current backend is TypeSafe's Jev model:
+
+| Function | Mode | Result |
+| --- | --- | --- |
+| `ai_match(text, prompt)` | Noul | `Float64` probability in **[0, 1]** that the prompt's statement is true |
+| `ai_choose(text, prompt, criteria)` | Choice | `String` containing the selected option's name |
+| `ai_score(text, prompt, criteria)` | Score | JSONB object containing `score`, `confidence`, and a `probabilities` array |
+
+Use scalar results or extracted JSON fields in `SELECT`, comparisons in `WHERE`,
+and `ORDER BY`.
+All arguments are SQL strings. Any SQL NULL argument produces SQL `NULL`
+without validating that row's criteria or calling the API.
+
+## Start GreptimeDB
+
+The functions are compiled and registered by the **default-enabled Cargo feature
+`ai_functions`**. Enable the Jev backend separately in the server process:
+
+```sh
+export GREPTIMEDB_EXPERIMENTAL_JEV=true
+export JEV_API_KEY=''
+cargo run -p cmd -- standalone start
+```
+
+If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first.
+The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to
+`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL).
+The MVP uses environment variables rather than TOML configuration.
+The Cargo feature includes all three SQL functions in the build; the runtime environment
+variables enable and configure its API calls. Default compilation does not turn
+on external API calls. To enable the feature explicitly, use
+`--features ai_functions`.
+
+## Query
+
+For a new example table, quote `service` and `message` in the DDL because the
+current parser treats them as keywords:
+
+```sql
+CREATE TABLE events (
+ occurred_at TIMESTAMP TIME INDEX,
+ "service" STRING,
+ "message" STRING,
+ PRIMARY KEY ("service")
+);
+```
+
+### Matching probability (Noul)
+
+```sql
+SELECT occurred_at, service, message
+FROM events
+WHERE occurred_at >= '2026-09-19T00:00:00Z'
+ AND occurred_at < '2026-09-20T00:00:00Z'
+ AND service = 'payments'
+ AND ai_match(
+ (message),
+ 'The event reports that a payment still failed after retries.'
+ ) >= 0.8
+ORDER BY occurred_at;
+```
+
+To return the scores and rank matching events:
+
+```sql
+SELECT occurred_at, message,
+ ai_match(message, 'The event reports that a payment still failed after retries.') AS score
+FROM events
+WHERE occurred_at >= '2026-09-19T00:00:00Z'
+ AND occurred_at < '2026-09-20T00:00:00Z'
+ AND service = 'payments'
+ORDER BY score DESC NULLS LAST;
+```
+
+Both arguments are strings; larger scores indicate a stronger match.
+`(message)` is an ordinary parenthesized string expression. For several columns,
+combine them explicitly, for example `concat(service, ': ', message)`.
+Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call.
+
+### Classification (Choice)
+
+`criteria` is a JSON object with **1 to 255 options**. Each key is an option name;
+its value is a description (string, object, or array), or JSON `null` if the name
+is sufficient. The function returns the selected key, suitable for filtering
+or grouping:
+
+```sql
+SELECT occurred_at, message,
+ ai_choose(message, 'Which team should handle this event?',
+ '{"billing":"Payments, invoices, refunds","technical":"Bugs and outages","other":null}') AS team
+FROM events
+WHERE occurred_at >= '2026-09-19T00:00:00Z'
+ AND occurred_at < '2026-09-20T00:00:00Z';
+```
+
+### Rating (Score)
+
+`criteria` is a JSON array of **2 to 10 ordered level descriptions**, from low
+to high. Descriptions can be strings, objects, or arrays. Level numbers start
+at zero. The function returns a JSONB object from one model evaluation:
+
+```json
+{
+ "score": 1.05,
+ "confidence": 0.92,
+ "probabilities": [0.0, 0.95, 0.05]
+}
+```
+
+- `score` is the provider's probability-weighted mean of the level numbers, in
+ `[0, N − 1]` for N levels. It may fall between levels.
+- `confidence` is the provider's confidence in `[0, 1]`, reflecting how concentrated
+ the distribution is. It is not a guarantee that the answer is correct.
+- `probabilities[i]` is the probability of `criteria[i]`, with exactly N entries in
+ level order. These are individual level probabilities, not cumulative probabilities.
+
+Different distributions can have the same score. `[0, 1, 0]` and `[0.5, 0, 0.5]`
+both have score `1`: the first concentrates on the middle level, while the second
+is split between the extremes. Inspect confidence and the distribution before
+treating a mean score as a definite severity level.
+
+For example, compute the rating once, extract numeric fields, and rank only rows
+meeting a chosen confidence threshold:
+
+```sql
+SELECT occurred_at, message,
+ json_get_float(rating, 'score') AS severity,
+ json_get_float(rating, 'confidence') AS confidence
+FROM (
+ SELECT occurred_at, message,
+ ai_score(message, 'How severe is this event?',
+ '["No impact to functionality","Degraded service with a workaround","Blocking issue with no workaround"]') AS rating
+ FROM events
+ WHERE occurred_at >= '2026-09-19T00:00:00Z'
+ AND occurred_at < '2026-09-20T00:00:00Z'
+) AS rated
+WHERE json_get_float(rating, 'confidence') >= 0.8
+ORDER BY severity DESC NULLS LAST;
+```
+
+Use `json_to_string(rating)` to display the full object, or
+`json_get_float(rating, 'probabilities[2]')` to read the probability of level 2.
+The array positions correspond to the supplied criteria, which provide the level
+descriptions.
+
+## Reusing an evaluation
+
+AI functions are marked volatile to prevent HTTP calls during planning. Identical
+calls written separately in `SELECT` and `WHERE` are evaluated separately. If the
+filter evaluates N non-null rows and M rows pass to the projection, this can cost
+N + M requests, reaching 2N when all rows pass.
+
+To return a result and filter by it, compute it once in a subquery and reference
+its alias in the outer query:
+
+```sql
+SELECT occurred_at, message, score
+FROM (
+ SELECT occurred_at, message,
+ ai_match(message, 'The event reports that a payment still failed after retries.') AS score
+ FROM events
+ WHERE occurred_at >= '2026-09-19T00:00:00Z'
+ AND occurred_at < '2026-09-20T00:00:00Z'
+ AND service = 'payments'
+) AS scored
+WHERE score >= 0.8
+ORDER BY score DESC;
+```
+
+This also applies to `ai_choose` and `ai_score`. Extract multiple JSON fields from
+the aliased `ai_score` result, as in the rating example above, rather than calling
+the model again for each field. Filter pushdown preserves the volatile projection
+instead of duplicating it into the outer predicate.
+
+## MVP behavior
+
+- Each non-null row makes one HTTP request per function invocation: the text is
+ `state`, and the prompt is the question's `instructions`. Choice and Score pass
+ the parsed JSON as `criteria`. Calls to different functions are separate requests.
+- Criteria for all non-null rows in a batch are validated before sending that
+ batch's requests. Malformed JSON, invalid description types, or invalid option/
+ level counts fail the query locally.
+- Answers must have the requested question type. Noul probabilities outside
+ `[0, 1]`, Choice labels not in the criteria, and Score values outside `[0, N − 1]`
+ fail the query rather than being clamped or replaced with NULL. Score responses
+ must also contain a numeric confidence in `[0, 1]` and probabilities for every
+ level. Each probability must be numeric and in `[0, 1]`; their sum must be within
+ `1e-6` of 1. Provider values are preserved without renormalizing the distribution
+ or recomputing the score.
+- Up to eight requests run concurrently per expression/batch invocation, not per
+ query or process. Concurrent partitions and queries can exceed eight requests
+ in total. Each request has a 30-second timeout. Errors (including rate limits
+ and invalid responses) fail the query. There is no automatic retry or
+ cross-query cache.
+- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the
+ candidate set, but SQL does not guarantee left-to-right predicate evaluation or
+ that `LIMIT` bounds the number of API calls.
+- Text passed to the function is sent to the configured TypeSafe endpoint.
+ The API key remains in the server environment, not in SQL.
+- This MVP is intended for standalone use. A distributed deployment would need
+ the environment on every node evaluating the function and separate validation.
+
+API reference:
+
+## Before stabilizing
+
+The experimental MVP still needs the following controls before stabilization:
+
+- A process-wide concurrency limit shared by AI function invocations, such as a semaphore.
+- Bounded retries with exponential backoff for HTTP `429` and `529`, following
+ TypeSafe's rate-limit guidance.
+- Request budgets and metrics for API calls, latency, retries, and rate-limit errors.
+
+These are follow-up work, not guarantees provided by the current implementation.
+
+## Validation
+
+Check both the default build (where all three functions are registered) and an isolated
+`common-function` build without default features. The regular AI function tests use a local
+HTTP server and need no API key:
+
+```sh
+cargo nextest run -p common-function
+cargo nextest run -p common-function --no-default-features
+cargo nextest run -p common-function --no-default-features --features ai_functions
+```
+
+The sqlness runner explicitly includes `ai_functions` when building its test binary:
+
+```sh
+cargo sqlness bare -t ai_functions
+```
+
+If using `--bins-dir`, provide a binary built with `ai_functions` (included by
+default). CI covers the AI-enabled path through the regular unit tests and
+sqlness tests; the feature-off checks above can be run locally when needed.
+
+An opt-in test calls the real service on synthetic payment events:
+
+```sh
+source ~/.zshrc
+GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \
+ -E 'test(test_ai_match_live)' --run-ignored ignored-only
+```
diff --git a/docs/how-to/jev.md b/docs/how-to/jev.md
deleted file mode 100644
index 2a21f085946..00000000000
--- a/docs/how-to/jev.md
+++ /dev/null
@@ -1,119 +0,0 @@
-# Natural-language filtering with Jev (experimental MVP)
-
-`jev(text, statement, threshold)` returns whether Jev judges the statement true
-for the text with probability **greater than or equal to** the threshold.
-It is an asynchronous SQL scalar function usable in `WHERE` and `SELECT`.
-
-## Start GreptimeDB
-
-Jev is compiled and registered by the **default-enabled Cargo feature
-`ai_functions`**. Enable API evaluation separately in the server process:
-
-```sh
-export GREPTIMEDB_EXPERIMENTAL_JEV=true
-export JEV_API_KEY=''
-cargo run -p cmd -- standalone start
-```
-
-If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first.
-The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to
-`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL).
-The MVP uses environment variables rather than TOML configuration.
-The Cargo feature includes the SQL function in the build; the runtime environment
-variables enable and configure its API calls. Default compilation does not turn
-on external API calls. To enable the feature explicitly, use
-`--features ai_functions`.
-
-## Query
-
-For a new example table, quote `service` and `message` in the DDL because the
-current parser treats them as keywords:
-
-```sql
-CREATE TABLE events (
- occurred_at TIMESTAMP TIME INDEX,
- "service" STRING,
- "message" STRING,
- PRIMARY KEY ("service")
-);
-```
-
-```sql
-SELECT occurred_at, service, message
-FROM events
-WHERE occurred_at >= '2026-09-19T00:00:00Z'
- AND occurred_at < '2026-09-20T00:00:00Z'
- AND service = 'payments'
- AND jev(
- (message),
- 'The event reports that a payment still failed after retries.',
- 0.8
- )
-ORDER BY occurred_at;
-```
-
-The first two arguments are strings and the third is a number in `[0, 1]`.
-`(message)` is an ordinary parenthesized string expression. For several columns,
-combine them explicitly, for example `concat(service, ': ', message)`.
-Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call.
-
-## MVP behavior
-
-- Each non-null row makes one HTTP request: the text is `state`, and the statement
- is a `noul` question's `instructions`. The returned `noul` probability is compared
- with the threshold locally; this is not the separate Choice/Score confidence.
-- Up to eight requests run concurrently per expression/batch invocation, not per
- query or process. Concurrent partitions and queries can exceed eight requests
- in total. Each request has a 30-second timeout. Errors (including rate limits
- and invalid responses) fail the query. There is no automatic retry or
- cross-query cache.
-- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the
- candidate set, but SQL does not guarantee left-to-right predicate evaluation or
- that `LIMIT` bounds the number of API calls.
-- Text passed to the function is sent to the configured TypeSafe endpoint.
- The API key remains in the server environment, not in SQL.
-- This MVP is intended for standalone use. A distributed deployment would need
- the environment on every node evaluating the function and separate validation.
-
-API reference:
-
-## Before stabilizing
-
-The experimental MVP still needs the following controls before stabilization:
-
-- A process-wide concurrency limit shared by Jev invocations, such as a semaphore.
-- Bounded retries with exponential backoff for HTTP `429` and `529`, following
- TypeSafe's rate-limit guidance.
-- Request budgets and metrics for API calls, latency, retries, and rate-limit errors.
-
-These are follow-up work, not guarantees provided by the current implementation.
-
-## Validation
-
-Check both the default build (where `jev` is registered) and an isolated
-`common-function` build without default features. Jev's regular tests use a local
-HTTP server and need no API key:
-
-```sh
-cargo nextest run -p common-function
-cargo nextest run -p common-function --no-default-features
-cargo nextest run -p common-function --no-default-features --features ai_functions
-```
-
-The sqlness runner explicitly includes `ai_functions` when building its test binary:
-
-```sh
-cargo sqlness bare -t jev
-```
-
-If using `--bins-dir`, provide a binary built with `ai_functions` (included by
-default). CI covers the AI-enabled path through the regular unit tests and
-sqlness tests; the feature-off checks above can be run locally when needed.
-
-An opt-in test calls the real service on synthetic payment events:
-
-```sh
-source ~/.zshrc
-GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \
- -E 'test(test_jev_live)' --run-ignored ignored-only
-```
diff --git a/src/cmd/src/bin/query_perf_fixture/case.rs b/src/cmd/src/bin/query_perf_fixture/case.rs
index 0f5239a6b51..271c3928e35 100644
--- a/src/cmd/src/bin/query_perf_fixture/case.rs
+++ b/src/cmd/src/bin/query_perf_fixture/case.rs
@@ -185,7 +185,7 @@ pub(super) fn default_max_concurrent_flushes() -> u64 {
256
}
pub(super) fn default_worker_channel_capacity() -> u64 {
- 65526
+ 65_536
}
pub(super) fn default_max_inflight_requests() -> u64 {
3000
diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs
index 67e2e573f38..1c4a83e1397 100644
--- a/src/cmd/tests/load_config_test.rs
+++ b/src/cmd/tests/load_config_test.rs
@@ -30,6 +30,7 @@ use datanode::config::{DatanodeOptions, RegionEngineConfig, StorageConfig};
use file_engine::config::EngineConfig as FileEngineConfig;
use flow::FlownodeOptions;
use frontend::frontend::FrontendOptions;
+use frontend::service_config::PendingRowsBatcherOptions;
use meta_client::MetaClientOptions;
use meta_srv::metasrv::MetasrvOptions;
use meta_srv::selector::SelectorType;
@@ -206,6 +207,10 @@ fn test_load_frontend_example_config() {
);
let expected = GreptimeOptions:: {
component: FrontendOptions {
+ pending_rows_batcher: PendingRowsBatcherOptions {
+ logical_table: Some(Default::default()),
+ ..Default::default()
+ },
default_timezone: Some("UTC".to_string()),
default_column_prefix: Some("greptime".to_string()),
auto_create_table: true,
@@ -389,6 +394,10 @@ fn test_load_standalone_example_config() {
);
let expected = GreptimeOptions:: {
component: StandaloneOptions {
+ pending_rows_batcher: PendingRowsBatcherOptions {
+ logical_table: Some(Default::default()),
+ ..Default::default()
+ },
default_timezone: Some("UTC".to_string()),
default_column_prefix: Some("greptime".to_string()),
auto_create_table: true,
diff --git a/src/common/event-recorder/src/recorder.rs b/src/common/event-recorder/src/recorder.rs
index b6c18ed08de..3ea5c7e55a8 100644
--- a/src/common/event-recorder/src/recorder.rs
+++ b/src/common/event-recorder/src/recorder.rs
@@ -291,17 +291,29 @@ pub struct EventRecorderOptions {
skip_serializing_if = "event_type_filter_is_all"
)]
pub event_types: EventTypeFilterRef,
+ /// How long buffered events may wait before reaching the event handler.
+ ///
+ /// Not a configuration option: it is skipped by serde so it never appears in
+ /// config files or `config/config.md`. Integration tests shorten it so they
+ /// do not pay a flush window per DDL they assert on.
+ #[serde(skip, default = "default_flush_interval")]
+ pub flush_interval: Duration,
}
fn default_events_table_ttl() -> Duration {
DEFAULT_EVENTS_TABLE_TTL
}
+fn default_flush_interval() -> Duration {
+ DEFAULT_FLUSH_INTERVAL_SECONDS
+}
+
impl Default for EventRecorderOptions {
fn default() -> Self {
Self {
ttl: DEFAULT_EVENTS_TABLE_TTL,
event_types: Arc::new(EventTypeFilter::All),
+ flush_interval: DEFAULT_FLUSH_INTERVAL_SECONDS,
}
}
}
@@ -321,13 +333,18 @@ pub struct EventRecorderImpl {
impl EventRecorderImpl {
pub fn new(event_handler: Box) -> Self {
- Self::with_event_type_filter(event_handler, Arc::new(EventTypeFilter::All))
+ Self::with_event_type_filter(
+ event_handler,
+ Arc::new(EventTypeFilter::All),
+ DEFAULT_FLUSH_INTERVAL_SECONDS,
+ )
}
/// Creates an event recorder with an event-type filter.
pub fn with_event_type_filter(
event_handler: Box,
event_types: EventTypeFilterRef,
+ flush_interval: Duration,
) -> Self {
let (tx, rx) = channel(DEFAULT_CHANNEL_SIZE);
let cancel_token = CancellationToken::new();
@@ -339,13 +356,9 @@ impl EventRecorderImpl {
cancel_token: cancel_token.clone(),
};
- let processor = EventProcessor::new(
- rx,
- event_handler,
- DEFAULT_FLUSH_INTERVAL_SECONDS,
- DEFAULT_MAX_RETRY_TIMES,
- )
- .with_cancel_token(cancel_token);
+ let processor =
+ EventProcessor::new(rx, event_handler, flush_interval, DEFAULT_MAX_RETRY_TIMES)
+ .with_cancel_token(cancel_token);
// Spawn a background task to process the events.
let handle = tokio::spawn(async move {
@@ -611,6 +624,7 @@ mod tests {
count: count.clone(),
}),
event_type_filter.clone(),
+ DEFAULT_FLUSH_INTERVAL_SECONDS,
);
assert!(Arc::ptr_eq(
diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs
index 037ca5b175a..6583f7ba607 100644
--- a/src/common/function/src/function_registry.rs
+++ b/src/common/function/src/function_registry.rs
@@ -28,14 +28,14 @@ use crate::aggrs::count_hash::CountHash;
use crate::aggrs::vector::VectorFunction as VectorAggrFunction;
use crate::function::{Function, FunctionRef};
use crate::function_factory::ScalarFunctionFactory;
+#[cfg(feature = "ai_functions")]
+use crate::scalars::ai;
use crate::scalars::anomaly::AnomalyFunction;
use crate::scalars::avg_calc::AvgCalcFunction;
use crate::scalars::date::DateFunction;
use crate::scalars::expression::ExpressionFunction;
use crate::scalars::hll_count::HllCalcFunction;
use crate::scalars::ip::IpFunctions;
-#[cfg(feature = "ai_functions")]
-use crate::scalars::jev::JevFunction;
use crate::scalars::json::JsonFunction;
use crate::scalars::matches::MatchesFunction;
use crate::scalars::matches_term::MatchesTermFunction;
@@ -229,7 +229,7 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(||
MatchesFunction::register(&function_registry);
MatchesTermFunction::register(&function_registry);
#[cfg(feature = "ai_functions")]
- JevFunction::register(&function_registry);
+ ai::register(&function_registry);
// System and administration functions
SystemFunction::register(&function_registry);
@@ -384,11 +384,17 @@ mod tests {
}
#[test]
- fn test_jev_registration_matches_feature() {
- assert_eq!(
- FUNCTION_REGISTRY.get_function("jev").is_some(),
- cfg!(feature = "ai_functions")
- );
+ fn test_ai_registration_matches_feature() {
+ for name in ["ai_match", "ai_choose", "ai_score"] {
+ assert_eq!(
+ FUNCTION_REGISTRY.get_function(name).is_some(),
+ cfg!(feature = "ai_functions"),
+ "{name}"
+ );
+ }
+ for name in ["jev", "jev_choice", "jev_score"] {
+ assert!(FUNCTION_REGISTRY.get_function(name).is_none(), "{name}");
+ }
}
#[test]
diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs
index 61fb7ec632c..4bc4c17f5bb 100644
--- a/src/common/function/src/scalars.rs
+++ b/src/common/function/src/scalars.rs
@@ -12,13 +12,13 @@
// See the License for the specific language governing permissions and
// limitations under the License.
+#[cfg(feature = "ai_functions")]
+pub(crate) mod ai;
pub mod anomaly;
pub(crate) mod date;
pub mod expression;
#[cfg(feature = "geo")]
pub mod geo;
-#[cfg(feature = "ai_functions")]
-pub(crate) mod jev;
pub mod json;
pub mod matches;
pub mod matches_term;
diff --git a/src/common/function/src/scalars/ai.rs b/src/common/function/src/scalars/ai.rs
new file mode 100644
index 00000000000..f2e17dffd36
--- /dev/null
+++ b/src/common/function/src/scalars/ai.rs
@@ -0,0 +1,454 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#[cfg(test)]
+mod tests;
+
+use std::fmt;
+use std::hash::Hash;
+use std::marker::PhantomData;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::{Array, StringViewArray, new_null_array};
+use arrow::datatypes::DataType;
+use async_trait::async_trait;
+use datafusion_common::cast::as_string_view_array;
+use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err, not_impl_err};
+use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+use futures::{StreamExt, TryStreamExt, stream};
+use reqwest::Client;
+use serde::Serialize;
+use serde_json::{Map, Value, json};
+
+use crate::function_factory::ScalarFunctionFactory;
+use crate::function_registry::FunctionRegistry;
+
+/// Registers the experimental AI matching, classification, and rating functions.
+pub(crate) fn register(registry: &FunctionRegistry) {
+ AiFunction::::register(registry);
+ AiFunction::::register(registry);
+ AiFunction::::register(registry);
+}
+
+/// Shared asynchronous execution for AI functions using the Jev backend.
+#[derive(PartialEq, Eq, Hash)]
+struct AiFunction {
+ signature: Signature,
+ question: PhantomData,
+ enabled: bool,
+ api_key: Option,
+ endpoint: String,
+ model: String,
+}
+
+struct AiRequest<'a, Q: AiQuestion> {
+ text: &'a str,
+ prompt: &'a str,
+ criteria: Arc,
+}
+
+/// Borrows string arguments without expanding constants to the batch size.
+enum StringArgument<'a> {
+ Scalar(Option<&'a str>),
+ Array(&'a StringViewArray),
+}
+
+impl<'a> StringArgument<'a> {
+ fn try_new(arg: &'a ColumnarValue) -> Result {
+ match arg {
+ ColumnarValue::Scalar(ScalarValue::Utf8View(value)) => {
+ Ok(Self::Scalar(value.as_deref()))
+ }
+ ColumnarValue::Array(array) => Ok(Self::Array(as_string_view_array(array)?)),
+ _ => exec_err!("AI functions expect Utf8View arguments"),
+ }
+ }
+
+ fn value(&self, row: usize) -> Option<&'a str> {
+ match self {
+ Self::Scalar(value) => *value,
+ Self::Array(array) => {
+ let array = *array;
+ array.is_valid(row).then(|| array.value(row))
+ }
+ }
+ }
+}
+
+impl AiFunction {
+ fn register(registry: &FunctionRegistry) {
+ registry.register(ScalarFunctionFactory {
+ name: Q::NAME.to_string(),
+ factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()),
+ });
+ }
+
+ async fn evaluate(&self, client: &Client, request: AiRequest<'_, Q>) -> Result {
+ let mut question = json!({ "type": Q::TYPE, "instructions": request.prompt });
+ if Q::ARG_COUNT == 3 {
+ question["criteria"] = json!(request.criteria.as_ref());
+ }
+ let response: Value = client
+ .post(&self.endpoint)
+ .json(&json!({
+ "model": self.model,
+ "state": request.text,
+ "questions": { "matches": question }
+ }))
+ .send()
+ .await
+ .and_then(|response| response.error_for_status())
+ .map_err(|e| exec_datafusion_err!("{} request failed: {e}", Q::NAME))?
+ .json()
+ .await
+ .map_err(|e| exec_datafusion_err!("{} response is not valid JSON: {e}", Q::NAME))?;
+
+ let answer = &response["answers"]["matches"];
+ if answer["type"].as_str() != Some(Q::TYPE) {
+ return exec_err!(
+ "{} response must contain answers.matches with type {}",
+ Q::NAME,
+ Q::TYPE
+ );
+ }
+ Q::parse_answer(answer, request.criteria.as_ref())
+ }
+
+ fn prepare_requests<'a>(
+ &self,
+ args: &'a [ColumnarValue],
+ number_rows: usize,
+ ) -> Result>>> {
+ if args.len() != Q::ARG_COUNT {
+ return exec_err!("{} requires {} arguments", Q::NAME, Q::ARG_COUNT);
+ }
+ let args = args
+ .iter()
+ .map(StringArgument::try_new)
+ .collect::>>()?;
+ let scalar_criteria = args[2..]
+ .iter()
+ .all(|arg| matches!(arg, StringArgument::Scalar(_)));
+ let mut shared_criteria: Option> = None;
+
+ // Validate all non-null rows before making any billable requests.
+ (0..number_rows)
+ .map(|row| {
+ let values: Option> = args.iter().map(|arg| arg.value(row)).collect();
+ values
+ .map(|values| {
+ // Initialize lazily so NULL rows never validate otherwise invalid criteria.
+ let criteria = match &shared_criteria {
+ Some(criteria) => Arc::clone(criteria),
+ None => {
+ let criteria = Arc::new(Q::parse_criteria(&values[2..])?);
+ if scalar_criteria {
+ shared_criteria = Some(Arc::clone(&criteria));
+ }
+ criteria
+ }
+ };
+ Ok(AiRequest {
+ text: values[0],
+ prompt: values[1],
+ criteria,
+ })
+ })
+ .transpose()
+ })
+ .collect()
+ }
+
+ fn client(&self) -> Result {
+ if !self.enabled {
+ return exec_err!(
+ "{} is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it",
+ Q::NAME
+ );
+ }
+ let key = self
+ .api_key
+ .as_deref()
+ .filter(|key| !key.trim().is_empty())
+ .ok_or_else(|| {
+ exec_datafusion_err!("{} requires the JEV_API_KEY environment variable", Q::NAME)
+ })?;
+ let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}"))
+ .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?;
+ authorization.set_sensitive(true);
+ let mut headers = reqwest::header::HeaderMap::new();
+ headers.insert(reqwest::header::AUTHORIZATION, authorization);
+ Client::builder()
+ .default_headers(headers)
+ .timeout(Duration::from_secs(30))
+ .build()
+ .map_err(|e| exec_datafusion_err!("failed to create {} HTTP client: {e}", Q::NAME))
+ }
+}
+
+impl Default for AiFunction {
+ fn default() -> Self {
+ Self {
+ // External model evaluations must not be constant-folded during planning.
+ signature: Signature::exact(
+ vec![DataType::Utf8View; Q::ARG_COUNT],
+ Volatility::Volatile,
+ ),
+ question: PhantomData,
+ enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"),
+ api_key: std::env::var("JEV_API_KEY").ok(),
+ endpoint: std::env::var("JEV_ENDPOINT")
+ .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()),
+ model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()),
+ }
+ }
+}
+
+impl fmt::Debug for AiFunction {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ // Credentials must never appear in plans or diagnostic output.
+ f.debug_struct("AiFunction")
+ .field("name", &Q::NAME)
+ .field("signature", &self.signature)
+ .field("enabled", &self.enabled)
+ .field("model", &self.model)
+ .finish_non_exhaustive()
+ }
+}
+
+impl ScalarUDFImpl for AiFunction {
+ fn name(&self) -> &str {
+ Q::NAME
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result {
+ Ok(Q::return_type())
+ }
+
+ fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result {
+ not_impl_err!("{} can only be called from async contexts", Q::NAME)
+ }
+}
+
+#[async_trait]
+impl AsyncScalarUDFImpl for AiFunction {
+ async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result {
+ let rows = self.prepare_requests(&args.args, args.number_rows)?;
+
+ if rows.iter().all(Option::is_none) {
+ return Ok(ColumnarValue::Array(new_null_array(
+ &Q::return_type(),
+ args.number_rows,
+ )));
+ }
+
+ // Reuse connections within the batch; bound in-flight requests and preserve row order.
+ let client = self.client()?;
+ // Collect futures to avoid async_trait's higher-ranked lifetime inference issue.
+ let requests: Vec<_> = rows
+ .into_iter()
+ .map(|row| {
+ let client = &client;
+ async move {
+ match row {
+ Some(request) => self.evaluate(client, request).await,
+ None => ScalarValue::try_from(&Q::return_type()),
+ }
+ }
+ })
+ .collect();
+ // This limit is per expression/batch invocation, not per query or process.
+ // Concurrent partitions and queries can each have their own in-flight requests.
+ let answers: Vec = stream::iter(requests).buffered(8).try_collect().await?;
+ Ok(ColumnarValue::Array(ScalarValue::iter_to_array(answers)?))
+ }
+}
+
+/// The request criteria and scalar answer contract for an AI question type.
+trait AiQuestion: fmt::Debug + Eq + Hash + Send + Sync + 'static {
+ type Criteria: Serialize + Send + Sync;
+
+ const NAME: &'static str;
+ const TYPE: &'static str;
+ const ARG_COUNT: usize;
+
+ fn return_type() -> DataType;
+ fn parse_criteria(args: &[&str]) -> Result;
+ fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result;
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct Noul;
+
+impl AiQuestion for Noul {
+ type Criteria = ();
+
+ const NAME: &'static str = "ai_match";
+ const TYPE: &'static str = "noul";
+ const ARG_COUNT: usize = 2;
+
+ fn return_type() -> DataType {
+ DataType::Float64
+ }
+
+ fn parse_criteria(_args: &[&str]) -> Result<()> {
+ Ok(())
+ }
+
+ fn parse_answer(answer: &Value, _criteria: &()) -> Result {
+ numeric_answer(Self::NAME, &answer["noul"], "noul", 1.0)
+ .map(|probability| ScalarValue::Float64(Some(probability)))
+ }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct Choice;
+
+impl AiQuestion for Choice {
+ type Criteria = Map;
+
+ const NAME: &'static str = "ai_choose";
+ const TYPE: &'static str = "choice";
+ const ARG_COUNT: usize = 3;
+
+ fn return_type() -> DataType {
+ DataType::Utf8
+ }
+
+ fn parse_criteria(args: &[&str]) -> Result {
+ let criteria: Value = serde_json::from_str(args[0])
+ .map_err(|e| exec_datafusion_err!("ai_choose criteria is not valid JSON: {e}"))?;
+ match criteria {
+ Value::Object(options)
+ if (1..=255).contains(&options.len())
+ && options.values().all(|v| v.is_null() || is_description(v)) =>
+ {
+ Ok(options)
+ }
+ _ => exec_err!(
+ "ai_choose criteria must be a JSON object with 1 to 255 options; descriptions must be strings, objects, arrays, or null"
+ ),
+ }
+ }
+
+ fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result {
+ let choice = answer["choice"]
+ .as_str()
+ .filter(|choice| criteria.contains_key(*choice))
+ .ok_or_else(|| {
+ exec_datafusion_err!("ai_choose response must contain a choice from the criteria")
+ })?;
+ Ok(ScalarValue::Utf8(Some(choice.to_string())))
+ }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct Score;
+
+impl AiQuestion for Score {
+ type Criteria = Vec;
+
+ const NAME: &'static str = "ai_score";
+ const TYPE: &'static str = "score";
+ const ARG_COUNT: usize = 3;
+
+ fn return_type() -> DataType {
+ DataType::BinaryView
+ }
+
+ fn parse_criteria(args: &[&str]) -> Result {
+ let criteria: Value = serde_json::from_str(args[0])
+ .map_err(|e| exec_datafusion_err!("ai_score criteria is not valid JSON: {e}"))?;
+ match criteria {
+ Value::Array(levels)
+ if (2..=10).contains(&levels.len()) && levels.iter().all(is_description) =>
+ {
+ Ok(levels)
+ }
+ _ => exec_err!(
+ "ai_score criteria must be a JSON array with 2 to 10 levels; descriptions must be strings, objects, or arrays"
+ ),
+ }
+ }
+
+ fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result {
+ let score = numeric_answer(
+ Self::NAME,
+ &answer["score"],
+ "score",
+ (criteria.len() - 1) as f64,
+ )?;
+ let confidence = numeric_answer(Self::NAME, &answer["confidence"], "confidence", 1.0)?;
+ let probabilities = score_probabilities(&answer["probabilities"], criteria.len())?;
+ let object = jsonb::Object::from([
+ ("score".to_string(), jsonb::Value::from(score)),
+ ("confidence".to_string(), jsonb::Value::from(confidence)),
+ (
+ "probabilities".to_string(),
+ jsonb::Value::Array(probabilities.into_iter().map(jsonb::Value::from).collect()),
+ ),
+ ]);
+ Ok(ScalarValue::BinaryView(Some(
+ jsonb::Value::Object(object).to_vec(),
+ )))
+ }
+}
+
+fn score_probabilities(distribution: &Value, level_count: usize) -> Result> {
+ if distribution
+ .as_object()
+ .is_none_or(|probabilities| probabilities.len() != level_count)
+ {
+ return exec_err!(
+ "ai_score response must contain probabilities for all {level_count} levels"
+ );
+ }
+ let probabilities = (0..level_count)
+ .map(|level| {
+ numeric_answer(
+ Score::NAME,
+ &distribution[level.to_string()],
+ "probability",
+ 1.0,
+ )
+ })
+ .collect::>>()?;
+ // Allow small rounding differences without renormalizing the provider's distribution.
+ if (probabilities.iter().sum::() - 1.0).abs() > 1e-6 {
+ return exec_err!("ai_score response probabilities must sum to 1 within 1e-6");
+ }
+ Ok(probabilities)
+}
+
+fn is_description(description: &Value) -> bool {
+ matches!(
+ description,
+ Value::String(_) | Value::Object(_) | Value::Array(_)
+ )
+}
+
+fn numeric_answer(name: &str, answer: &Value, field: &str, max: f64) -> Result {
+ answer
+ .as_f64()
+ .filter(|score| (0.0..=max).contains(score))
+ .ok_or_else(|| {
+ exec_datafusion_err!("{name} response must contain a finite {field} in [0, {max}]")
+ })
+}
diff --git a/src/common/function/src/scalars/ai/tests.rs b/src/common/function/src/scalars/ai/tests.rs
new file mode 100644
index 00000000000..0b001e22e7a
--- /dev/null
+++ b/src/common/function/src/scalars/ai/tests.rs
@@ -0,0 +1,766 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use axum::http::{HeaderMap, StatusCode};
+use axum::routing::post;
+use axum::{Json, Router};
+use datafusion::assert_batches_eq;
+use datafusion::prelude::{SessionConfig, SessionContext};
+use tokio::net::TcpListener;
+use tokio::task::JoinHandle;
+
+use super::*;
+use crate::function::FunctionContext;
+use crate::function_registry::FUNCTION_REGISTRY;
+
+struct MockServer {
+ endpoint: String,
+ task: JoinHandle<()>,
+}
+
+impl MockServer {
+ async fn start(router: Router) -> Self {
+ let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap());
+ let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
+ Self { endpoint, task }
+ }
+
+ fn function(&self) -> AiFunction {
+ AiFunction {
+ enabled: true,
+ api_key: Some("test-key".to_string()),
+ endpoint: self.endpoint.clone(),
+ model: "test-model".to_string(),
+ ..Default::default()
+ }
+ }
+}
+
+impl Drop for MockServer {
+ fn drop(&mut self) {
+ self.task.abort();
+ }
+}
+
+fn context(function: AiFunction) -> SessionContext {
+ let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));
+ ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf());
+ let json_get_float = FUNCTION_REGISTRY.get_function("json_get_float").unwrap();
+ ctx.register_udf(json_get_float.provide(FunctionContext::default()));
+ ctx
+}
+
+#[test]
+fn test_ai_scalar_criteria_share_allocation_and_skip_null_rows() {
+ fn assert_shared_criteria(criteria: &str) {
+ let function = AiFunction::::default();
+ for (texts, criteria) in [
+ (vec![None, Some("first"), Some("second")], criteria),
+ (vec![None, None], "not json"),
+ (vec![], "not json"),
+ ] {
+ let number_rows = texts.len();
+ let expected_requests = texts.iter().flatten().count();
+ let args = vec![
+ ColumnarValue::Array(Arc::new(StringViewArray::from(texts))),
+ ColumnarValue::Scalar(ScalarValue::Utf8View(Some("prompt".to_string()))),
+ ColumnarValue::Scalar(ScalarValue::Utf8View(Some(criteria.to_string()))),
+ ];
+ let requests = function.prepare_requests(&args, number_rows).unwrap();
+ assert_eq!(requests.len(), number_rows);
+ assert_eq!(requests.iter().flatten().count(), expected_requests);
+
+ // Regression: parsed constant criteria must not be duplicated for each row.
+ let mut criteria = requests.iter().flatten().map(|request| &request.criteria);
+ if let Some(first) = criteria.next() {
+ assert!(criteria.all(|other| Arc::ptr_eq(first, other)));
+ }
+ }
+ }
+
+ assert_shared_criteria::(r#"{"billing":null,"technical":"errors"}"#);
+ assert_shared_criteria::(r#"["low","high"]"#);
+}
+
+#[tokio::test]
+async fn test_ai_match_sql_returns_and_filters_noul_probability() {
+ let server = MockServer::start(Router::new().route(
+ "/v1/systemone",
+ post(
+ |headers: HeaderMap, Json(request): Json| async move {
+ assert_eq!(headers["authorization"], "Bearer test-key");
+ assert_eq!(request["model"], "test-model");
+ assert_eq!(request["questions"]["matches"]["type"], "noul");
+ // The prompt is passed through verbatim.
+ assert_eq!(
+ request["questions"]["matches"]["instructions"],
+ "payment failed"
+ );
+ let probability = match request["state"].as_str().unwrap() {
+ "failed" => {
+ // Complete this row after later rows to exercise result alignment.
+ tokio::time::sleep(Duration::from_millis(20)).await;
+ 0.95
+ }
+ "boundary" => 0.8,
+ "recovered" => 0.1,
+ "impossible" => 0.0,
+ "certain" => 1.0,
+ unexpected => panic!("unexpected state: {unexpected}"),
+ };
+ Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}}))
+ },
+ ),
+ ))
+ .await;
+ let ctx = context(server.function::());
+ let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL), (5, 'impossible'), (6, 'certain')) AS events(id, message)";
+ let batches = ctx
+ .sql(&format!(
+ "SELECT id FROM {events} WHERE ai_match((message), 'payment failed') >= 0.8 ORDER BY id"
+ ))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_batches_eq!(
+ [
+ "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 6 |", "+----+"
+ ],
+ &batches
+ );
+
+ let batches = ctx
+ .sql(&format!(
+ "SELECT id, ai_match(message, 'payment failed') AS score FROM {events} ORDER BY score DESC NULLS LAST"
+ ))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Float64);
+ assert_batches_eq!(
+ [
+ "+----+-------+",
+ "| id | score |",
+ "+----+-------+",
+ "| 6 | 1.0 |",
+ "| 1 | 0.95 |",
+ "| 2 | 0.8 |",
+ "| 3 | 0.1 |",
+ "| 5 | 0.0 |",
+ "| 4 | |",
+ "+----+-------+"
+ ],
+ &batches
+ );
+}
+
+#[tokio::test]
+async fn test_ai_match_nulls_need_no_api_key() {
+ let ctx = context(AiFunction:: {
+ enabled: true,
+ api_key: None,
+ ..Default::default()
+ });
+ let batches = ctx
+ .sql("SELECT ai_match(NULL, 'condition') AS a, ai_match('text', NULL) AS b")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(batches[0].schema().field(0).data_type(), &DataType::Float64);
+ assert_batches_eq!(
+ [
+ "+---+---+",
+ "| a | b |",
+ "+---+---+",
+ "| | |",
+ "+---+---+"
+ ],
+ &batches
+ );
+
+ let error = ctx
+ .sql("SELECT ai_match('text', 'condition')")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("JEV_API_KEY"), "{error}");
+
+ let ctx = context(AiFunction:: {
+ enabled: false,
+ api_key: None,
+ ..Default::default()
+ });
+ let error = ctx
+ .sql("SELECT ai_match('text', 'condition')")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap_err();
+ assert!(
+ error
+ .to_string()
+ .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"),
+ "{error}"
+ );
+}
+
+#[tokio::test]
+async fn test_ai_match_bad_api_responses_fail_the_query() {
+ for (status, body, expected) in [
+ (
+ StatusCode::TOO_MANY_REQUESTS,
+ "{}",
+ "ai_match request failed",
+ ),
+ (
+ StatusCode::OK,
+ "not json",
+ "ai_match response is not valid JSON",
+ ),
+ (
+ StatusCode::OK,
+ r#"{"answers":{}}"#,
+ "ai_match response must contain",
+ ),
+ (
+ StatusCode::OK,
+ r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#,
+ "ai_match response must contain",
+ ),
+ (
+ StatusCode::OK,
+ r#"{"answers":{"matches":{"type":"noul","noul":-0.1}}}"#,
+ "ai_match response must contain",
+ ),
+ (
+ StatusCode::OK,
+ r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#,
+ "ai_match response must contain",
+ ),
+ ] {
+ let server = MockServer::start(
+ Router::new().route("/v1/systemone", post(move || async move { (status, body) })),
+ )
+ .await;
+ let ctx = context(server.function::());
+ let error = ctx
+ .sql("SELECT ai_match('text', 'condition')")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains(expected), "{error}");
+ }
+}
+
+#[tokio::test]
+async fn test_ai_choose_returns_and_filters_supplied_labels() {
+ let server = MockServer::start(Router::new().route(
+ "/v1/systemone",
+ post(|Json(request): Json| async move {
+ let question = &request["questions"]["matches"];
+ assert_eq!(question["type"], "choice");
+ assert_eq!(question["instructions"], "Route the ticket");
+ let choice = match request["state"].as_str().unwrap() {
+ "refund" => {
+ assert_eq!(question["criteria"], json!({"billing": null, "technical": {"examples": ["crash"]}}));
+ tokio::time::sleep(Duration::from_millis(20)).await;
+ "billing"
+ }
+ "crash" => {
+ assert_eq!(question["criteria"], json!({"billing": ["refunds"], "technical": "errors"}));
+ "technical"
+ }
+ unexpected => panic!("unexpected state: {unexpected}"),
+ };
+ Json(json!({"answers": {"matches": {"type": "choice", "choice": choice, "confidence": 0.7}}}))
+ }),
+ ))
+ .await;
+ let ctx = context(server.function::());
+ let events = r#"(VALUES
+ (1, 'refund', 'Route the ticket', '{"billing":null,"technical":{"examples":["crash"]}}'),
+ (2, 'crash', 'Route the ticket', '{"billing":["refunds"],"technical":"errors"}'),
+ (3, NULL, 'Route the ticket', 'invalid JSON'),
+ (4, 'refund', NULL, 'invalid JSON'),
+ (5, 'refund', 'Route the ticket', NULL)
+ ) AS events(id, message, prompt, criteria)"#;
+ let batches = ctx
+ .sql(&format!(
+ "SELECT id, ai_choose(message, prompt, criteria) AS team FROM {events} ORDER BY id"
+ ))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Utf8);
+ assert_batches_eq!(
+ [
+ "+----+-----------+",
+ "| id | team |",
+ "+----+-----------+",
+ "| 1 | billing |",
+ "| 2 | technical |",
+ "| 3 | |",
+ "| 4 | |",
+ "| 5 | |",
+ "+----+-----------+"
+ ],
+ &batches
+ );
+ let batches = ctx
+ .sql(&format!(
+ "SELECT id FROM {events} WHERE ai_choose(message, prompt, criteria) = 'billing'"
+ ))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches);
+
+ let batches = ctx
+ .sql(
+ r#"SELECT id, ai_choose(message, 'Route the ticket',
+ '{"billing":null,"technical":{"examples":["crash"]}}') AS team
+ FROM (VALUES (1, 'refund'), (2, NULL), (3, 'refund')) AS events(id, message)
+ ORDER BY id"#,
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_batches_eq!(
+ [
+ "+----+---------+",
+ "| id | team |",
+ "+----+---------+",
+ "| 1 | billing |",
+ "| 2 | |",
+ "| 3 | billing |",
+ "+----+---------+"
+ ],
+ &batches
+ );
+}
+
+#[tokio::test]
+async fn test_ai_score_returns_fractional_scores_on_each_rows_scale() {
+ let server = MockServer::start(Router::new().route(
+ "/v1/systemone",
+ post(|Json(request): Json| async move {
+ let question = &request["questions"]["matches"];
+ assert_eq!(question["type"], "score");
+ assert_eq!(question["instructions"], "Rate severity");
+ let (score, probabilities) = match request["state"].as_str().unwrap() {
+ "low" => {
+ assert_eq!(question["criteria"], json!(["low", {"description": "medium"}, ["high"]]));
+ tokio::time::sleep(Duration::from_millis(20)).await;
+ (0.0, json!({"0":1.0,"1":0.0,"2":0.0}))
+ }
+ "fractional" => (1.25, json!({"2":0.25,"0":0.0,"1":0.75})),
+ "high" => (2.0, json!({"0":0.0,"1":0.0,"2":1.0})),
+ "binary" => {
+ assert_eq!(question["criteria"], json!(["low", "high"]));
+ (1.0, json!({"0":0.0,"1":1.0}))
+ }
+ "ten levels" => {
+ assert_eq!(question["criteria"].as_array().unwrap().len(), 10);
+ let probabilities: Map = (0..10)
+ .map(|level| (level.to_string(), json!(if level == 9 { 1.0 } else { 0.0 })))
+ .collect();
+ (9.0, Value::Object(probabilities))
+ }
+ unexpected => panic!("unexpected state: {unexpected}"),
+ };
+ Json(json!({"answers": {"matches": {"type": "score", "score": score, "confidence": 0.4, "probabilities": probabilities}}}))
+ }),
+ ))
+ .await;
+ let ctx = context(server.function::());
+ let batches = ctx
+ .sql(
+ r#"
+ SELECT id, json_get_float(ai_score(message, 'Rate severity', criteria), 'score') AS severity
+ FROM (VALUES
+ (1, 'low', '["low", {"description":"medium"}, ["high"]]'),
+ (2, 'fractional', '["low", "medium", "high"]'),
+ (3, 'high', '["low", "medium", "high"]'),
+ (4, 'binary', '["low", "high"]'),
+ (5, 'ten levels', '["a","b","c","d","e","f","g","h","i","j"]'),
+ (6, NULL, 'invalid JSON'),
+ (7, 'high', NULL)
+ ) AS events(id, message, criteria)
+ ORDER BY severity DESC NULLS LAST, id
+ "#,
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Float64);
+ assert_batches_eq!(
+ [
+ "+----+----------+",
+ "| id | severity |",
+ "+----+----------+",
+ "| 5 | 9.0 |",
+ "| 3 | 2.0 |",
+ "| 2 | 1.25 |",
+ "| 4 | 1.0 |",
+ "| 1 | 0.0 |",
+ "| 6 | |",
+ "| 7 | |",
+ "+----+----------+"
+ ],
+ &batches
+ );
+
+ let batches = ctx
+ .sql(
+ r#"SELECT id, json_get_float(ai_score(message, 'Rate severity',
+ '["low",{"description":"medium"},["high"]]'), 'score') AS severity
+ FROM (VALUES (1, 'low'), (2, 'fractional'), (3, 'high'), (4, NULL)) AS events(id, message)
+ ORDER BY id"#,
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_batches_eq!(
+ [
+ "+----+----------+",
+ "| id | severity |",
+ "+----+----------+",
+ "| 1 | 0.0 |",
+ "| 2 | 1.25 |",
+ "| 3 | 2.0 |",
+ "| 4 | |",
+ "+----+----------+"
+ ],
+ &batches
+ );
+}
+
+#[tokio::test]
+async fn test_ai_score_preserves_uncertainty_in_one_evaluation() {
+ let calls = Arc::new(AtomicUsize::new(0));
+ let request_calls = calls.clone();
+ let server = MockServer::start(Router::new().route(
+ "/v1/systemone",
+ post(move |Json(request): Json| {
+ request_calls.fetch_add(1, Ordering::Relaxed);
+ async move {
+ let (confidence, probabilities) = match request["state"].as_str().unwrap() {
+ "medium" => (1.0, json!({"2":0.0,"0":0.0,"1":1.0})),
+ "split" => (0.25, json!({"2":0.5,"0":0.5,"1":0.0})),
+ unexpected => panic!("unexpected state: {unexpected}"),
+ };
+ Json(json!({"answers":{"matches":{
+ "type":"score", "score":1.0, "confidence":confidence,
+ "probabilities":probabilities
+ }}}))
+ }
+ }),
+ ))
+ .await;
+ let ctx = context(server.function::());
+ let rated = r#"(SELECT id, ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating
+ FROM (VALUES (1, 'medium'), (2, 'split'), (3, NULL)) AS events(id, message)) AS rated"#;
+ let batches = ctx
+ .sql(&format!(
+ "SELECT id, json_get_float(rating, 'score') AS score,
+ json_get_float(rating, 'confidence') AS confidence,
+ json_get_float(rating, 'probabilities[0]') AS p0,
+ json_get_float(rating, 'probabilities[1]') AS p1,
+ json_get_float(rating, 'probabilities[2]') AS p2
+ FROM {rated} ORDER BY id"
+ ))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_batches_eq!(
+ [
+ "+----+-------+------------+-----+-----+-----+",
+ "| id | score | confidence | p0 | p1 | p2 |",
+ "+----+-------+------------+-----+-----+-----+",
+ "| 1 | 1.0 | 1.0 | 0.0 | 1.0 | 0.0 |",
+ "| 2 | 1.0 | 0.25 | 0.5 | 0.0 | 0.5 |",
+ "| 3 | | | | | |",
+ "+----+-------+------------+-----+-----+-----+"
+ ],
+ &batches
+ );
+ assert_eq!(calls.load(Ordering::Relaxed), 2);
+
+ let batches = ctx
+ .sql(&format!(
+ "SELECT id FROM {rated} WHERE json_get_float(rating, 'confidence') >= 0.8
+ ORDER BY json_get_float(rating, 'score') DESC"
+ ))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches);
+ assert_eq!(calls.load(Ordering::Relaxed), 4);
+}
+
+#[test]
+fn test_ai_score_preserves_provider_precision() {
+ let answer = json!({
+ "score":1.6, "confidence":0.92,
+ "probabilities":{"2":0.6999995,"0":0.1,"1":0.2}
+ });
+ let criteria = vec![json!("low"), json!("medium"), json!("high")];
+ let ScalarValue::BinaryView(Some(bytes)) = Score::parse_answer(&answer, &criteria).unwrap()
+ else {
+ panic!("expected a JSONB object");
+ };
+ let object: Value =
+ serde_json::from_str(&jsonb::from_slice(&bytes).unwrap().to_string()).unwrap();
+ assert_eq!(object.as_object().unwrap().len(), 3);
+ assert_eq!(object["score"].as_f64(), Some(1.6));
+ assert_eq!(object["confidence"].as_f64(), Some(0.92));
+ let probabilities: Vec = object["probabilities"]
+ .as_array()
+ .unwrap()
+ .iter()
+ .map(|probability| probability.as_f64().unwrap())
+ .collect();
+ assert_eq!(probabilities, vec![0.1, 0.2, 0.6999995]);
+}
+
+#[tokio::test]
+async fn test_ai_score_rejects_invalid_uncertainty() {
+ for (field, value) in [
+ ("confidence", Value::Null),
+ ("confidence", json!(-0.1)),
+ ("confidence", json!(1.1)),
+ ("confidence", json!("0.9")),
+ ("probabilities", Value::Null),
+ ("probabilities", json!([0.0, 1.0, 0.0])),
+ ("probabilities", json!({"0":1.0})),
+ ("probabilities", json!({"0":0.5,"2":0.5,"3":0.0})),
+ ("probabilities", json!({"0":0.0,"1":1.0,"2":0.0,"3":0.0})),
+ ("probabilities", json!({"0":-0.1,"1":0.6,"2":0.5})),
+ ("probabilities", json!({"0":0.0,"1":1.1,"2":0.0})),
+ ("probabilities", json!({"0":0.0,"1":0.95,"2":"0.05"})),
+ ("probabilities", json!({"0":0.0,"1":0.0,"2":0.0})),
+ ("probabilities", json!({"0":0.5,"1":0.5,"2":0.5})),
+ ] {
+ let mut answer = json!({
+ "type":"score", "score":1.0, "confidence":0.9,
+ "probabilities":{"0":0.0,"1":1.0,"2":0.0}
+ });
+ answer[field] = value;
+ let server = MockServer::start(Router::new().route(
+ "/v1/systemone",
+ post(move || {
+ let answer = answer.clone();
+ async move { Json(json!({"answers":{"matches":answer}})) }
+ }),
+ ))
+ .await;
+ let ctx = context(server.function::());
+ let error = ctx
+ .sql(r#"SELECT ai_score('text', 'prompt', '["low","medium","high"]')"#)
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("ai_score response"), "{error}");
+ }
+}
+
+#[tokio::test]
+async fn test_ai_invalid_criteria_fail_before_requests() {
+ let choice_ctx = context(AiFunction:: {
+ enabled: true,
+ api_key: None,
+ ..Default::default()
+ });
+ let score_ctx = context(AiFunction:: {
+ enabled: true,
+ api_key: None,
+ ..Default::default()
+ });
+ let too_many_choices: Map = (0..256)
+ .map(|index| (index.to_string(), Value::Null))
+ .collect();
+ for (ctx, name, valid, invalid) in [
+ (
+ &choice_ctx,
+ "ai_choose",
+ r#"{"billing":null}"#,
+ vec![
+ "not json".to_string(),
+ "null".to_string(),
+ "{}".to_string(),
+ "[]".to_string(),
+ r#"{"billing":1}"#.to_string(),
+ r#"{"billing":true}"#.to_string(),
+ Value::Object(too_many_choices).to_string(),
+ ],
+ ),
+ (
+ &score_ctx,
+ "ai_score",
+ r#"["low","high"]"#,
+ vec![
+ "not json".to_string(),
+ "null".to_string(),
+ "{}".to_string(),
+ "[]".to_string(),
+ r#"["low"]"#.to_string(),
+ r#"["low",null]"#.to_string(),
+ r#"["low",1]"#.to_string(),
+ r#"["low",true]"#.to_string(),
+ json!(vec!["level"; 11]).to_string(),
+ ],
+ ),
+ ] {
+ let error = ctx
+ .sql(&format!(
+ "SELECT {name}(message, 'prompt', 'not json') FROM (VALUES (NULL), ('text')) AS input(message)"
+ ))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap_err();
+ assert!(
+ error.to_string().contains(&format!("{name} criteria")),
+ "{error}"
+ );
+ for criteria in invalid {
+ // A later invalid row must fail before even creating a client for the first row.
+ let error = ctx.sql(&format!(
+ "SELECT {name}('text', 'prompt', criteria) FROM (VALUES ('{valid}'), ('{criteria}')) AS input(criteria)"
+ )).await.unwrap().collect().await.unwrap_err();
+ assert!(
+ error.to_string().contains(&format!("{name} criteria")),
+ "{error}"
+ );
+ }
+ }
+}
+
+#[tokio::test]
+async fn test_ai_choose_and_score_reject_invalid_answers() {
+ for (name, criteria, answer) in [
+ (
+ "ai_choose",
+ r#"{"billing":null}"#,
+ json!({"type":"score","choice":"billing"}),
+ ),
+ (
+ "ai_choose",
+ r#"{"billing":null}"#,
+ json!({"type":"choice","choice":"unknown"}),
+ ),
+ (
+ "ai_choose",
+ r#"{"billing":null}"#,
+ json!({"type":"choice","choice":1}),
+ ),
+ ("ai_choose", r#"{"billing":null}"#, json!({"type":"choice"})),
+ (
+ "ai_score",
+ r#"["low","medium","high"]"#,
+ json!({"type":"noul","score":0.5}),
+ ),
+ (
+ "ai_score",
+ r#"["low","medium","high"]"#,
+ json!({"type":"score","score":2.1}),
+ ),
+ (
+ "ai_score",
+ r#"["low","high"]"#,
+ json!({"type":"score","score":-0.1}),
+ ),
+ (
+ "ai_score",
+ r#"["low","high"]"#,
+ json!({"type":"score","score":"NaN"}),
+ ),
+ ("ai_score", r#"["low","high"]"#, json!({"type":"score"})),
+ ] {
+ let server = MockServer::start(Router::new().route(
+ "/v1/systemone",
+ post(move || {
+ let answer = answer.clone();
+ async move { Json(json!({"answers":{"matches":answer}})) }
+ }),
+ ))
+ .await;
+ let ctx = context(server.function::());
+ ctx.register_udf(
+ AsyncScalarUDF::new(Arc::new(server.function::())).into_scalar_udf(),
+ );
+ let error = ctx
+ .sql(&format!("SELECT {name}('text', 'prompt', '{criteria}')"))
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap_err();
+ assert!(
+ error
+ .to_string()
+ .contains(&format!("{name} response must contain")),
+ "{error}"
+ );
+ }
+}
+
+#[tokio::test]
+#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"]
+async fn test_ai_match_live() {
+ let ctx = SessionContext::new();
+ let function = FUNCTION_REGISTRY.get_function("ai_match").unwrap();
+ ctx.register_udf(function.provide(FunctionContext::default()));
+ let batches = ctx.sql(
+ "SELECT id FROM (VALUES
+ (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'),
+ (2, 'Payment succeeded on the second retry; the payment is complete.'),
+ (3, 'User logged in successfully.')
+ ) AS events(id, message)
+ WHERE ai_match(message, 'The event reports that a payment still failed after retries.') >= 0.8
+ ORDER BY id",
+ ).await.unwrap().collect().await.unwrap();
+ assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches);
+}
diff --git a/src/common/function/src/scalars/jev.rs b/src/common/function/src/scalars/jev.rs
deleted file mode 100644
index 1c6475e6fac..00000000000
--- a/src/common/function/src/scalars/jev.rs
+++ /dev/null
@@ -1,209 +0,0 @@
-// Copyright 2023 Greptime Team
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#[cfg(test)]
-mod tests;
-
-use std::fmt;
-use std::sync::Arc;
-use std::time::Duration;
-
-use arrow::array::BooleanArray;
-use arrow::datatypes::DataType;
-use async_trait::async_trait;
-use datafusion_common::cast::{as_float64_array, as_string_view_array};
-use datafusion_common::utils::take_function_args;
-use datafusion_common::{Result, exec_datafusion_err, exec_err, not_impl_err};
-use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
-use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
-use futures::{StreamExt, TryStreamExt, stream};
-use reqwest::Client;
-use serde_json::{Value, json};
-
-use crate::function_factory::ScalarFunctionFactory;
-use crate::function_registry::FunctionRegistry;
-
-/// Experimental natural-language predicate: `jev(text, statement, threshold)`.
-/// Each non-null row is a Noul question; its probability is compared with `>=`.
-#[derive(PartialEq, Eq, Hash)]
-pub(crate) struct JevFunction {
- signature: Signature,
- enabled: bool,
- api_key: Option,
- endpoint: String,
- model: String,
-}
-
-impl JevFunction {
- pub(crate) fn register(registry: &FunctionRegistry) {
- registry.register(ScalarFunctionFactory {
- name: "jev".to_string(),
- factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()),
- });
- }
-
- async fn evaluate(&self, client: &Client, text: &str, statement: &str) -> Result {
- let response: Value = client
- .post(&self.endpoint)
- .json(&json!({
- "model": self.model,
- "state": text,
- "questions": {
- "matches": { "type": "noul", "instructions": statement }
- }
- }))
- .send()
- .await
- .and_then(|response| response.error_for_status())
- .map_err(|e| exec_datafusion_err!("jev request failed: {e}"))?
- .json()
- .await
- .map_err(|e| exec_datafusion_err!("jev response is not valid JSON: {e}"))?;
-
- let answer = &response["answers"]["matches"];
- let probability = answer["noul"].as_f64().filter(|p| (0.0..=1.0).contains(p));
- match (answer["type"].as_str(), probability) {
- (Some("noul"), Some(probability)) => Ok(probability),
- _ => exec_err!(
- "jev response must contain answers.matches with type noul and a probability in [0, 1]"
- ),
- }
- }
-
- fn client(&self) -> Result {
- if !self.enabled {
- return exec_err!(
- "jev is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it"
- );
- }
- let key = self
- .api_key
- .as_deref()
- .filter(|key| !key.trim().is_empty())
- .ok_or_else(|| {
- exec_datafusion_err!("jev requires the JEV_API_KEY environment variable")
- })?;
- let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}"))
- .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?;
- authorization.set_sensitive(true);
- let mut headers = reqwest::header::HeaderMap::new();
- headers.insert(reqwest::header::AUTHORIZATION, authorization);
- Client::builder()
- .default_headers(headers)
- .timeout(Duration::from_secs(30))
- .build()
- .map_err(|e| exec_datafusion_err!("failed to create jev HTTP client: {e}"))
- }
-}
-
-impl Default for JevFunction {
- fn default() -> Self {
- Self {
- // External model evaluations must not be constant-folded during planning.
- signature: Signature::exact(
- vec![DataType::Utf8View, DataType::Utf8View, DataType::Float64],
- Volatility::Volatile,
- ),
- enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"),
- api_key: std::env::var("JEV_API_KEY").ok(),
- endpoint: std::env::var("JEV_ENDPOINT")
- .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()),
- model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()),
- }
- }
-}
-
-impl fmt::Debug for JevFunction {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- // Credentials must never appear in plans or diagnostic output.
- f.debug_struct("JevFunction")
- .field("signature", &self.signature)
- .field("enabled", &self.enabled)
- .field("model", &self.model)
- .finish_non_exhaustive()
- }
-}
-
-impl ScalarUDFImpl for JevFunction {
- fn name(&self) -> &str {
- "jev"
- }
-
- fn signature(&self) -> &Signature {
- &self.signature
- }
-
- fn return_type(&self, _arg_types: &[DataType]) -> Result {
- Ok(DataType::Boolean)
- }
-
- fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result {
- not_impl_err!("jev can only be called from async contexts")
- }
-}
-
-#[async_trait]
-impl AsyncScalarUDFImpl for JevFunction {
- async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result {
- let arrays = args
- .args
- .into_iter()
- .map(|arg| arg.into_array(args.number_rows))
- .collect::>>()?;
- let [texts, statements, thresholds] = take_function_args(self.name(), arrays)?;
- let texts = as_string_view_array(&texts)?;
- let statements = as_string_view_array(&statements)?;
- let thresholds = as_float64_array(&thresholds)?;
- let rows: Vec<_> = texts
- .iter()
- .zip(statements.iter())
- .zip(thresholds.iter())
- .map(|((text, statement), threshold)| Some((text?, statement?, threshold?)))
- .collect();
-
- // Validate the whole batch before making any billable requests.
- for (_, _, threshold) in rows.iter().flatten() {
- if !(0.0..=1.0).contains(threshold) {
- return exec_err!("jev threshold must be finite and in [0, 1]");
- }
- }
- if rows.iter().all(Option::is_none) {
- return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null(
- args.number_rows,
- ))));
- }
-
- // Reuse connections within the batch; bound in-flight requests and preserve row order.
- let client = self.client()?;
- let requests: Vec<_> = rows
- .into_iter()
- .map(|row| {
- let client = &client;
- async move {
- match row {
- Some((text, statement, threshold)) => self
- .evaluate(client, text, statement)
- .await
- .map(|p| Some(p >= threshold)),
- None => Ok(None),
- }
- }
- })
- .collect();
- // This limit is per expression/batch invocation, not per query or process.
- // Concurrent partitions and queries can each have their own in-flight requests.
- let matches: Vec> = stream::iter(requests).buffered(8).try_collect().await?;
- Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(matches))))
- }
-}
diff --git a/src/common/function/src/scalars/jev/tests.rs b/src/common/function/src/scalars/jev/tests.rs
deleted file mode 100644
index 82d7ab03013..00000000000
--- a/src/common/function/src/scalars/jev/tests.rs
+++ /dev/null
@@ -1,250 +0,0 @@
-// Copyright 2023 Greptime Team
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use axum::http::{HeaderMap, StatusCode};
-use axum::routing::post;
-use axum::{Json, Router};
-use datafusion::assert_batches_eq;
-use datafusion::prelude::{SessionConfig, SessionContext};
-use tokio::net::TcpListener;
-use tokio::task::JoinHandle;
-
-use super::*;
-use crate::function::FunctionContext;
-use crate::function_registry::FUNCTION_REGISTRY;
-
-struct MockServer {
- endpoint: String,
- task: JoinHandle<()>,
-}
-
-impl MockServer {
- async fn start(router: Router) -> Self {
- let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
- let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap());
- let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
- Self { endpoint, task }
- }
-
- fn function(&self) -> JevFunction {
- JevFunction {
- enabled: true,
- api_key: Some("test-key".to_string()),
- endpoint: self.endpoint.clone(),
- model: "test-model".to_string(),
- ..Default::default()
- }
- }
-}
-
-impl Drop for MockServer {
- fn drop(&mut self) {
- self.task.abort();
- }
-}
-
-fn context(function: JevFunction) -> SessionContext {
- let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));
- ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf());
- ctx
-}
-
-#[tokio::test]
-async fn test_jev_sql_filters_using_noul_probability() {
- let server = MockServer::start(Router::new().route(
- "/v1/systemone",
- post(
- |headers: HeaderMap, Json(request): Json| async move {
- assert_eq!(headers["authorization"], "Bearer test-key");
- assert_eq!(request["model"], "test-model");
- assert_eq!(request["questions"]["matches"]["type"], "noul");
- // The SQL statement is passed through verbatim, not rewritten as a prompt.
- assert_eq!(
- request["questions"]["matches"]["instructions"],
- "payment failed"
- );
- let probability = match request["state"].as_str().unwrap() {
- "failed" => {
- // Complete this row after later rows to exercise result alignment.
- tokio::time::sleep(Duration::from_millis(20)).await;
- 0.95
- }
- "boundary" => 0.8,
- "recovered" => 0.1,
- unexpected => panic!("unexpected state: {unexpected}"),
- };
- Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}}))
- },
- ),
- ))
- .await;
- let ctx = context(server.function());
- let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL)) AS events(id, message)";
- let batches = ctx
- .sql(&format!(
- "SELECT id FROM {events} WHERE jev((message), 'payment failed', 0.8) ORDER BY id"
- ))
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
- assert_batches_eq!(
- ["+----+", "| id |", "+----+", "| 1 |", "| 2 |", "+----+"],
- &batches
- );
-
- let batches = ctx
- .sql(&format!(
- "SELECT id, jev(message, 'payment failed', 0.8) AS matched FROM {events} ORDER BY id"
- ))
- .await
- .unwrap()
- .collect()
- .await
- .unwrap();
- assert_batches_eq!(
- [
- "+----+---------+",
- "| id | matched |",
- "+----+---------+",
- "| 1 | true |",
- "| 2 | true |",
- "| 3 | false |",
- "| 4 | |",
- "+----+---------+"
- ],
- &batches
- );
-}
-
-#[tokio::test]
-async fn test_jev_nulls_and_invalid_thresholds_need_no_api_key() {
- let ctx = context(JevFunction {
- enabled: true,
- api_key: None,
- ..Default::default()
- });
- let batches = ctx
- .sql("SELECT jev(NULL, 'condition', 0.8) AS a, jev('text', NULL, 0.8) AS b, jev('text', 'condition', NULL) AS c")
- .await.unwrap().collect().await.unwrap();
- assert_batches_eq!(
- [
- "+---+---+---+",
- "| a | b | c |",
- "+---+---+---+",
- "| | | |",
- "+---+---+---+"
- ],
- &batches
- );
-
- for threshold in ["-0.1", "1.1", "'NaN'::DOUBLE", "'Infinity'::DOUBLE"] {
- let error = ctx
- .sql(&format!("SELECT jev('text', 'condition', {threshold})"))
- .await
- .unwrap()
- .collect()
- .await
- .unwrap_err();
- assert!(
- error
- .to_string()
- .contains("jev threshold must be finite and in [0, 1]"),
- "{error}"
- );
- }
- let error = ctx
- .sql("SELECT jev('text', 'condition', 0.8)")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap_err();
- assert!(error.to_string().contains("JEV_API_KEY"), "{error}");
-
- let ctx = context(JevFunction {
- enabled: false,
- api_key: None,
- ..Default::default()
- });
- let error = ctx
- .sql("SELECT jev('text', 'condition', 0.8)")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap_err();
- assert!(
- error
- .to_string()
- .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"),
- "{error}"
- );
-}
-
-#[tokio::test]
-async fn test_jev_bad_api_responses_fail_the_query() {
- for (status, body, expected) in [
- (StatusCode::TOO_MANY_REQUESTS, "{}", "jev request failed"),
- (StatusCode::OK, "not json", "jev response is not valid JSON"),
- (
- StatusCode::OK,
- r#"{"answers":{}}"#,
- "jev response must contain",
- ),
- (
- StatusCode::OK,
- r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#,
- "jev response must contain",
- ),
- (
- StatusCode::OK,
- r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#,
- "jev response must contain",
- ),
- ] {
- let server = MockServer::start(
- Router::new().route("/v1/systemone", post(move || async move { (status, body) })),
- )
- .await;
- let ctx = context(server.function());
- let error = ctx
- .sql("SELECT jev('text', 'condition', 0.8)")
- .await
- .unwrap()
- .collect()
- .await
- .unwrap_err();
- assert!(error.to_string().contains(expected), "{error}");
- }
-}
-
-#[tokio::test]
-#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"]
-async fn test_jev_live() {
- let ctx = SessionContext::new();
- let function = FUNCTION_REGISTRY.get_function("jev").unwrap();
- ctx.register_udf(function.provide(FunctionContext::default()));
- let batches = ctx.sql(
- "SELECT id FROM (VALUES
- (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'),
- (2, 'Payment succeeded on the second retry; the payment is complete.'),
- (3, 'User logged in successfully.')
- ) AS events(id, message)
- WHERE jev(message, 'The event reports that a payment still failed after retries.', 0.8)
- ORDER BY id",
- ).await.unwrap().collect().await.unwrap();
- assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches);
-}
diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs
index d5214ae5ddb..ffa148c7d0b 100644
--- a/src/common/grpc-expr/src/alter.rs
+++ b/src/common/grpc-expr/src/alter.rs
@@ -24,21 +24,23 @@ use api::v1::{
SkippingIndexType as PbSkippingIndexType, column_def,
};
use common_query::AddColumnLocation;
+use datatypes::json::{JsonSettings, JsonTypeHint};
+use datatypes::prelude::ConcreteDataType;
use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions};
use snafu::{OptionExt, ResultExt, ensure};
use store_api::region_request::{SetRegionOption, UnsetRegionOption};
use table::metadata::{TableId, TableMeta};
use table::requests::{
AddColumnRequest, AlterKind, AlterTableRequest, AnnotationFamily, ModifyColumnTypeRequest,
- SetDefaultRequest, SetIndexOption, UnsetIndexOption,
+ SetDefaultRequest, SetIndexOption, SetJsonSettingsRequest, UnsetIndexOption,
};
use crate::error::{
self, ColumnNotFoundSnafu, InvalidColumnDefSnafu, InvalidIndexOptionSnafu,
- InvalidSetFulltextOptionRequestSnafu, InvalidSetSkippingIndexOptionRequestSnafu,
- InvalidSetTableOptionRequestSnafu, InvalidUnsetTableOptionRequestSnafu,
- MissingAlterIndexOptionSnafu, MissingFieldSnafu, MissingTableMetaSnafu,
- MissingTimestampColumnSnafu, Result, UnknownLocationTypeSnafu,
+ InvalidJsonSettingsSnafu, InvalidSetFulltextOptionRequestSnafu,
+ InvalidSetSkippingIndexOptionRequestSnafu, InvalidSetTableOptionRequestSnafu,
+ InvalidUnsetTableOptionRequestSnafu, MissingAlterIndexOptionSnafu, MissingFieldSnafu,
+ MissingTableMetaSnafu, MissingTimestampColumnSnafu, Result, UnknownLocationTypeSnafu,
};
const LOCATION_TYPE_FIRST: i32 = LocationType::First as i32;
@@ -59,6 +61,34 @@ fn annotation_family_of_keys<'a>(
})
}
+fn json_settings_from_proto(settings: api::v1::JsonSettings) -> Result {
+ let type_hints = settings
+ .type_hints
+ .into_iter()
+ .map(|hint| {
+ let data_type = ConcreteDataType::from(
+ ColumnDataTypeWrapper::try_new(hint.data_type, hint.datatype_extension)
+ .context(error::ColumnDataTypeSnafu)?,
+ );
+
+ Ok(JsonTypeHint {
+ path: hint.path,
+ data_type,
+ // Index configuration is not supported yet, so this is temporarily
+ // hardcoded to false.
+ inverted_index: false,
+ })
+ })
+ .collect::>>()?;
+
+ JsonSettings::try_new(type_hints, settings.max_auto_expanded_paths).map_err(|err| {
+ InvalidJsonSettingsSnafu {
+ err: err.to_string(),
+ }
+ .build()
+ })
+}
+
/// Returns the annotation family when `kind` is a SET/UNSET whose keys all
/// belong to one family — the alters that only rewrite table metadata and skip
/// region dispatch. A mixed batch is an error; never interpret it as "not an
@@ -189,6 +219,17 @@ pub fn alter_expr_to_request(
columns: modify_column_type_requests,
}
}
+ Kind::SetJsonSettings(set_json_settings) => {
+ let settings = set_json_settings
+ .settings
+ .context(MissingFieldSnafu { field: "settings" })?;
+ AlterKind::SetJsonSettings {
+ request: SetJsonSettingsRequest {
+ column_name: set_json_settings.column_name,
+ settings: json_settings_from_proto(settings)?,
+ },
+ }
+ }
Kind::DropColumns(DropColumns { drop_columns }) => AlterKind::DropColumns {
names: drop_columns.into_iter().map(|c| c.name).collect(),
},
diff --git a/src/common/grpc-expr/src/error.rs b/src/common/grpc-expr/src/error.rs
index ccc40efb5cb..6c9302c39f4 100644
--- a/src/common/grpc-expr/src/error.rs
+++ b/src/common/grpc-expr/src/error.rs
@@ -73,6 +73,13 @@ pub enum Error {
source: api::error::Error,
},
+ #[snafu(display("Invalid JSON settings: {}", err))]
+ InvalidJsonSettings {
+ err: String,
+ #[snafu(implicit)]
+ location: Location,
+ },
+
#[snafu(display("Unknown location type: {}", location_type))]
UnknownLocationType {
location_type: i32,
@@ -188,6 +195,7 @@ impl ErrorExt for Error {
| Error::MissingTimestampColumn { .. } => StatusCode::InvalidArguments,
Error::MissingField { .. } => StatusCode::InvalidArguments,
Error::InvalidColumnDef { source, .. } => source.status_code(),
+ Error::InvalidJsonSettings { .. } => StatusCode::InvalidArguments,
Error::UnknownLocationType { .. } => StatusCode::InvalidArguments,
Error::UnknownColumnDataType { .. } | Error::InvalidFulltextIndexColumnType { .. } => {
diff --git a/src/common/meta/src/ddl/alter_table/executor.rs b/src/common/meta/src/ddl/alter_table/executor.rs
index 53bc4b3d05b..419aa9ae12d 100644
--- a/src/common/meta/src/ddl/alter_table/executor.rs
+++ b/src/common/meta/src/ddl/alter_table/executor.rs
@@ -333,6 +333,7 @@ fn build_new_table_info(
}
AlterKind::DropColumns { .. }
| AlterKind::ModifyColumnTypes { .. }
+ | AlterKind::SetJsonSettings { .. }
| AlterKind::SetTableOptions { .. }
| AlterKind::UnsetTableOptions { .. }
| AlterKind::SetAnnotations { .. }
diff --git a/src/common/meta/src/ddl/alter_table/region_request.rs b/src/common/meta/src/ddl/alter_table/region_request.rs
index 73edd4c3615..b4f7b679422 100644
--- a/src/common/meta/src/ddl/alter_table/region_request.rs
+++ b/src/common/meta/src/ddl/alter_table/region_request.rs
@@ -90,6 +90,7 @@ fn create_proto_alter_kind(
})))
}
Kind::ModifyColumnTypes(x) => Ok(Some(alter_request::Kind::ModifyColumnTypes(x.clone()))),
+ Kind::SetJsonSettings(x) => Ok(Some(alter_request::Kind::SetJsonSettings(x.clone()))),
Kind::DropColumns(x) => {
let drop_columns = x
.drop_columns
diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs
index 670f3c7c1dc..c1ab0e5e5a2 100644
--- a/src/common/meta/src/ddl/event/table.rs
+++ b/src/common/meta/src/ddl/event/table.rs
@@ -193,6 +193,7 @@ pub(crate) fn alter_table_kind_name(kind: &AlterTableKind) -> Option<&'static st
AlterTableKind::DropColumns(_) => Some("drop_columns"),
AlterTableKind::RenameTable(_) => Some("rename_table"),
AlterTableKind::ModifyColumnTypes(_) => Some("modify_column_types"),
+ AlterTableKind::SetJsonSettings(_) => Some("set_json_settings"),
AlterTableKind::SetTableOptions(_) => Some("set_table_options"),
AlterTableKind::UnsetTableOptions(_) => Some("unset_table_options"),
AlterTableKind::SetIndex(_) => Some("set_index"),
diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs
index f2ed05534d6..8ad388c2ff7 100644
--- a/src/datatypes/src/data_type.rs
+++ b/src/datatypes/src/data_type.rs
@@ -492,7 +492,9 @@ impl TryFrom<&ArrowDataType> for ConcreteDataType {
ArrowDataType::Decimal128(precision, scale) => {
ConcreteDataType::decimal128_datatype(*precision, *scale)
}
- ArrowDataType::Struct(fields) => ConcreteDataType::Struct(StructType::from(fields)),
+ ArrowDataType::Struct(fields) => {
+ ConcreteDataType::Struct(StructType::try_from_arrow_fields(fields)?)
+ }
ArrowDataType::Float16
| ArrowDataType::Date64
| ArrowDataType::FixedSizeBinary(_)
diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs
index 3aa0a4c9792..2e86d7791c5 100644
--- a/src/datatypes/src/extension/json.rs
+++ b/src/datatypes/src/extension/json.rs
@@ -24,9 +24,10 @@ use parquet_variant_compute::VariantType;
use serde::{Deserialize, Serialize};
use snafu::{ResultExt, ensure};
-use crate::error::InvalidJson2LayoutSnafu;
+use crate::error::{InvalidJson2LayoutSnafu, SerializeSnafu};
pub use crate::json::JSON2_REMAINDER_FIELD_NAME;
use crate::json::JsonSettings;
+use crate::schema::Metadata;
/// Aligns JSON2 field types with their built arrays while preserving field and schema metadata.
pub fn align_schema_with_json_array(schema: SchemaRef, columns: &[ArrayRef]) -> SchemaRef {
@@ -151,7 +152,7 @@ pub struct JsonMetadata {
}
impl JsonMetadata {
- /// Creates metadata for the JSON2 layout version 2.
+ /// Creates metadata for the latest JSON2 layout (currently V2).
pub fn new(json_settings: JsonSettings) -> Self {
Self {
json_settings,
@@ -300,6 +301,27 @@ impl ExtensionType for Json2ExtensionType {
}
}
+/// Returns JSON2 column metadata with updated settings and the latest layout.
+///
+/// Existing SSTs retain their own physical layout metadata.
+pub fn json2_metadata_with_updated_settings(
+ current_metadata: &Metadata,
+ settings: JsonSettings,
+) -> crate::error::Result {
+ let json_metadata = JsonMetadata::new(settings);
+
+ let mut metadata = current_metadata.clone();
+ metadata.insert(
+ EXTENSION_TYPE_NAME_KEY.to_string(),
+ Json2ExtensionType::NAME.to_string(),
+ );
+ metadata.insert(
+ EXTENSION_TYPE_METADATA_KEY.to_string(),
+ serde_json::to_string(&json_metadata).context(SerializeSnafu)?,
+ );
+ Ok(metadata)
+}
+
/// Checks whether this field is either a legacy JSONB or JSON2 extension type.
pub fn is_any_json_extension_type>(field: T) -> bool {
let name = field.as_ref().extension_type_name();
@@ -442,6 +464,35 @@ mod tests {
Ok(())
}
+ #[test]
+ fn test_json2_metadata_with_updated_settings_upgrades_layout() -> crate::error::Result<()> {
+ for (name, json) in [
+ (JsonExtensionType::NAME, r#"{"json_settings":{}}"#),
+ (Json2ExtensionType::NAME, r#"{"json_settings":{}}"#),
+ (
+ Json2ExtensionType::NAME,
+ r#"{"json_settings":{},"layout_version":2}"#,
+ ),
+ ] {
+ let metadata = HashMap::from([
+ (EXTENSION_TYPE_NAME_KEY.to_string(), name.to_string()),
+ (EXTENSION_TYPE_METADATA_KEY.to_string(), json.to_string()),
+ ("other".to_string(), "kept".to_string()),
+ ]);
+ let settings = JsonSettings::try_new(vec![], Some(10))?;
+ let updated = json2_metadata_with_updated_settings(&metadata, settings.clone())?;
+ assert_eq!(Some("kept"), updated.get("other").map(String::as_str));
+ assert_eq!(
+ Some(Json2ExtensionType::NAME),
+ updated.get(EXTENSION_TYPE_NAME_KEY).map(String::as_str)
+ );
+ let json_metadata: JsonMetadata =
+ serde_json::from_str(updated.get(EXTENSION_TYPE_METADATA_KEY).unwrap()).unwrap();
+ assert_eq!(JsonMetadata::new(settings), json_metadata);
+ }
+ Ok(())
+ }
+
#[test]
fn test_parse_json2_physical_layout() -> crate::error::Result<()> {
let legacy = Field::new("data", DataType::Struct(Fields::empty()), true)
diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs
index a5b4940e048..cd73e202bf3 100644
--- a/src/datatypes/src/json.rs
+++ b/src/datatypes/src/json.rs
@@ -23,6 +23,7 @@ pub mod value;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
+use std::sync::Arc;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value as Json};
@@ -31,6 +32,7 @@ use snafu::ResultExt;
use crate::data_type::ConcreteDataType;
use crate::error::{self, InvalidJson2SettingsSnafu, Result, UnsupportedJsonTypeSnafu};
use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb};
+use crate::prelude::DataType as _;
use crate::types::json_type::{JsonNativeType, JsonObjectType};
use crate::value::{ListValue, StructValue, Value};
@@ -93,6 +95,15 @@ pub struct JsonTypeHint {
pub inverted_index: bool,
}
+/// Specifies how encoding handles a value that does not match its JSON2 type hint.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum TypeHintMismatchPolicy {
+ /// Returns an error for a type hint mismatch.
+ Reject,
+ /// Coerces to the hint type, encoding unconvertible values as null.
+ CoerceOrNull,
+}
+
/// Context for JSON encoding/decoding that tracks the current key path.
#[derive(Clone, Debug)]
pub struct JsonContext<'a> {
@@ -135,6 +146,30 @@ impl JsonSettings {
(self.type_hints, self.max_auto_expanded_paths)
}
+ /// Returns whether the settings are equivalent, ignoring type hint order.
+ pub fn equivalent(&self, other: &Self) -> bool {
+ let Self {
+ type_hints,
+ max_auto_expanded_paths,
+ } = self;
+ let Self {
+ type_hints: other_hints,
+ max_auto_expanded_paths: other_max_auto_expanded_paths,
+ } = other;
+
+ if max_auto_expanded_paths != other_max_auto_expanded_paths
+ || type_hints.len() != other_hints.len()
+ {
+ return false;
+ }
+
+ let mut hints = type_hints.iter().collect::>();
+ let mut other_hints = other_hints.iter().collect::>();
+ hints.sort_unstable_by(|a, b| a.path.cmp(&b.path));
+ other_hints.sort_unstable_by(|a, b| a.path.cmp(&b.path));
+ hints == other_hints
+ }
+
/// Decode an encoded StructValue back into a serde_json::Value.
pub fn decode(&self, value: Value) -> Result {
let mut context = JsonContext {
@@ -146,6 +181,15 @@ impl JsonSettings {
/// Encode a serde_json::Value into a Value::Json using current settings.
pub fn encode(&self, json: Json) -> Result {
+ self.encode_with_type_hint_mismatch_policy(json, TypeHintMismatchPolicy::Reject)
+ }
+
+ /// Encodes a serde_json::Value using the given type hint mismatch policy.
+ pub fn encode_with_type_hint_mismatch_policy(
+ &self,
+ json: Json,
+ policy: TypeHintMismatchPolicy,
+ ) -> Result {
if let Json::Object(object) = &json
&& object.contains_key(JSON2_REMAINDER_FIELD_NAME)
{
@@ -160,7 +204,7 @@ impl JsonSettings {
path: Vec::new(),
settings: self,
};
- encode_json_with_context(json, &mut context).map(|v| Value::Json(Box::new(v)))
+ encode_json_with_context(json, &mut context, policy).map(|v| Value::Json(Box::new(v)))
}
}
@@ -272,36 +316,41 @@ fn with_key_context(
}
/// Main encoding function with key path tracking
-fn encode_json_with_context(json: Json, context: &mut JsonContext) -> Result {
+fn encode_json_with_context(
+ json: Json,
+ context: &mut JsonContext,
+ policy: TypeHintMismatchPolicy,
+) -> Result {
if context.path.is_empty() && !matches!(json, Json::Object(_)) {
return UnsupportedJsonTypeSnafu.fail();
}
match json {
- Json::Object(json_object) => encode_json_object_with_context(json_object, context),
- Json::Array(json_array) => encode_json_array_with_context(json_array, context),
- _ => encode_json_value_with_context(json, context),
+ Json::Object(json_object) => encode_json_object_with_context(json_object, context, policy),
+ Json::Array(json_array) => encode_json_array_with_context(json_array, context, policy),
+ _ => encode_json_value_with_context(json, context, policy),
}
}
fn encode_json_object_with_context<'a>(
json_object: Map,
context: &mut JsonContext<'a>,
+ policy: TypeHintMismatchPolicy,
) -> Result {
let mut object = BTreeMap::new();
for (key, value) in json_object {
let value = with_key_context(context, &key, |context| {
if let Some(hint) = context.type_hint() {
- encode_json_value_with_hint(value, hint, context)
+ encode_json_value_with_hint(value, hint, context, policy)
} else {
- encode_json_value_with_context(value, context)
+ encode_json_value_with_context(value, context, policy)
}
})?;
object.insert(key, value.into_variant());
}
- fill_missing_type_hints(&mut object, context)?;
+ fill_missing_type_hints(&mut object, context, policy)?;
Ok(JsonValue::new(JsonVariant::Object(object)))
}
@@ -309,13 +358,14 @@ fn encode_json_object_with_context<'a>(
fn fill_missing_type_hints(
object: &mut BTreeMap,
context: &mut JsonContext,
+ policy: TypeHintMismatchPolicy,
) -> Result<()> {
for hint in &context.settings.type_hints {
if hint.path.len() > context.path.len() && hint.path.starts_with(&context.path) {
let depth = context.path.len();
let key = &hint.path[depth];
with_key_context(context, key, |context| {
- insert_missing_type_hint(object, context, hint, depth)
+ insert_missing_type_hint(object, context, hint, depth, policy)
})?;
}
}
@@ -327,6 +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 {
if json.is_null() {
return Ok(JsonValue::null());
}
- let invalid_type = || {
- error::InvalidJsonSnafu {
+ let json = match (&hint.data_type, json) {
+ (ConcreteDataType::String(_), Json::String(value)) => {
+ return Ok(value.into());
+ }
+ (
+ ConcreteDataType::Int8(_)
+ | ConcreteDataType::Int16(_)
+ | ConcreteDataType::Int32(_)
+ | ConcreteDataType::Int64(_),
+ Json::Number(value),
+ ) => {
+ if let Some(value) = value.as_i64() {
+ return Ok(value.into());
+ }
+ Json::Number(value)
+ }
+ (
+ ConcreteDataType::UInt8(_)
+ | ConcreteDataType::UInt16(_)
+ | ConcreteDataType::UInt32(_)
+ | ConcreteDataType::UInt64(_),
+ Json::Number(value),
+ ) => {
+ if let Some(value) = value.as_u64() {
+ return Ok(value.into());
+ }
+ Json::Number(value)
+ }
+ (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), Json::Number(value)) => {
+ if let Some(value) = value.as_f64() {
+ return Ok(value.into());
+ }
+ Json::Number(value)
+ }
+ (ConcreteDataType::Boolean(_), Json::Bool(value)) => {
+ return Ok(value.into());
+ }
+ (_, json) => json,
+ };
+
+ match policy {
+ TypeHintMismatchPolicy::Reject => error::InvalidJsonSnafu {
value: format!(
"JSON value at {} does not match JSON2 type hint {}",
context.path.join("."),
hint.data_type
),
}
- .fail()
- };
+ .fail(),
+ TypeHintMismatchPolicy::CoerceOrNull => {
+ let value = coerce_json_value_to_type(json, &hint.data_type);
+ let value = Json::try_from(value).map_err(|error| {
+ error::InvalidJsonSnafu {
+ value: error.to_string(),
+ }
+ .build()
+ })?;
+ Ok(JsonValue::new(value.into()))
+ }
+ }
+}
- match (&hint.data_type, json) {
- (ConcreteDataType::String(_), Json::String(v)) => Ok(v.into()),
- (
- ConcreteDataType::Int8(_)
- | ConcreteDataType::Int16(_)
- | ConcreteDataType::Int32(_)
- | ConcreteDataType::Int64(_),
- Json::Number(v),
- ) => match v.as_i64() {
- Some(v) => Ok(v.into()),
- None => invalid_type(),
- },
- (
- ConcreteDataType::UInt8(_)
- | ConcreteDataType::UInt16(_)
- | ConcreteDataType::UInt32(_)
- | ConcreteDataType::UInt64(_),
- Json::Number(v),
- ) => match v.as_u64() {
- Some(v) => Ok(v.into()),
- None => invalid_type(),
- },
- (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), Json::Number(v)) => {
- match v.as_f64() {
- Some(v) => Ok(v.into()),
- None => invalid_type(),
+/// Coerces a JSON value using the same semantics as JSON2 query projection.
+/// Values that cannot be converted become null.
+pub(crate) fn coerce_json_value_to_type(value: Json, to_type: &ConcreteDataType) -> Value {
+ if value.is_null() {
+ return Value::Null;
+ }
+
+ if to_type.is_string() {
+ let value = match value {
+ Json::String(value) => value,
+ value => value.to_string(),
+ };
+ return Value::String(value.into());
+ }
+
+ if matches!(to_type, ConcreteDataType::Binary(_)) {
+ return Value::Binary(encode_serde_json_as_jsonb(value).into());
+ }
+
+ if let Some(struct_type) = to_type.as_struct() {
+ let Json::Object(mut object) = value else {
+ return Value::Null;
+ };
+ let values = struct_type
+ .fields()
+ .iter()
+ .map(|field| {
+ object
+ .remove(field.name())
+ .map(|value| coerce_json_value_to_type(value, field.data_type()))
+ .unwrap_or(Value::Null)
+ })
+ .collect::>();
+ return Value::Struct(StructValue::new(values, struct_type.clone()));
+ }
+
+ if let Some(list_type) = to_type.as_list() {
+ let Json::Array(values) = value else {
+ return Value::Null;
+ };
+ let item_type = list_type.item_type().clone();
+ let values = values
+ .into_iter()
+ .map(|value| coerce_json_value_to_type(value, &item_type))
+ .collect::>();
+ return Value::List(ListValue::new(values, Arc::new(item_type)));
+ }
+
+ let value = match value {
+ Json::Bool(value) => Value::Boolean(value),
+ Json::Number(value) => {
+ if let Some(value) = value.as_i64() {
+ Value::Int64(value)
+ } else if let Some(value) = value.as_u64() {
+ Value::UInt64(value)
+ } else if let Some(value) = value.as_f64() {
+ Value::Float64(value.into())
+ } else {
+ Value::Null
}
}
- (ConcreteDataType::Boolean(_), Json::Bool(v)) => Ok(v.into()),
- _ => invalid_type(),
- }
+ Json::String(value) => Value::String(value.into()),
+ Json::Array(_) | Json::Object(_) | Json::Null => Value::Null,
+ };
+ to_type.try_cast(value).unwrap_or(Value::Null)
}
fn encode_json_array_with_context<'a>(
json_array: Vec,
context: &mut JsonContext<'a>,
+ policy: TypeHintMismatchPolicy,
) -> Result {
let json_array_len = json_array.len();
let mut items = Vec::with_capacity(json_array_len);
for (index, value) in json_array.into_iter().enumerate() {
let item_value = with_key_context(context, &index.to_string(), |context| {
- encode_json_value_with_context(value, context)
+ encode_json_value_with_context(value, context, policy)
})?;
items.push(item_value);
}
@@ -458,7 +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 {
+fn encode_json_value_with_context(
+ json: Json,
+ context: &mut JsonContext,
+ policy: TypeHintMismatchPolicy,
+) -> Result {
if context.path.len() >= JSON2_MAX_STRUCTURED_DEPTH
&& matches!(&json, Json::Object(_) | Json::Array(_))
{
@@ -487,8 +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> {
@@ -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();
diff --git a/src/datatypes/src/types/struct_type.rs b/src/datatypes/src/types/struct_type.rs
index ffd1d1c0b25..fdbfd4886b9 100644
--- a/src/datatypes/src/types/struct_type.rs
+++ b/src/datatypes/src/types/struct_type.rs
@@ -19,6 +19,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Field};
use arrow_schema::Fields;
use serde::{Deserialize, Serialize};
+use crate::error::Result;
use crate::prelude::{ConcreteDataType, DataType, LogicalTypeId};
use crate::value::Value;
use crate::vectors::StructVectorBuilder;
@@ -100,6 +101,24 @@ impl StructType {
}
}
+ /// Like the `From<&Fields>` impl, but returns an error for arrow field
+ /// types greptimedb cannot represent instead of panicking.
+ pub fn try_from_arrow_fields(fields: &Fields) -> Result {
+ let struct_fields = fields
+ .iter()
+ .map(|field| {
+ Ok(StructField::new(
+ field.name().clone(),
+ ConcreteDataType::try_from(field.data_type())?,
+ field.is_nullable(),
+ ))
+ })
+ .collect::>>()?;
+ Ok(StructType {
+ fields: Arc::new(struct_fields),
+ })
+ }
+
pub fn fields(&self) -> Arc> {
self.fields.clone()
}
diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs
index 56a551484f2..da1df882aba 100644
--- a/src/datatypes/src/value.rs
+++ b/src/datatypes/src/value.rs
@@ -1152,12 +1152,13 @@ impl TryFrom for Value {
ScalarValue::UInt16(u) => Value::from(u),
ScalarValue::UInt32(u) => Value::from(u),
ScalarValue::UInt64(u) => Value::from(u),
- ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) => {
+ ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s) => {
Value::from(s.map(StringBytes::from))
}
ScalarValue::Binary(b)
| ScalarValue::LargeBinary(b)
- | ScalarValue::FixedSizeBinary(_, b) => Value::from(b.map(Bytes::from)),
+ | ScalarValue::FixedSizeBinary(_, b)
+ | ScalarValue::BinaryView(b) => Value::from(b.map(Bytes::from)),
ScalarValue::List(array) => {
// this is for item type
let datatype = ConcreteDataType::try_from(&array.value_type())?;
@@ -1221,7 +1222,15 @@ impl TryFrom for Value {
.map(|v| Value::Decimal128(Decimal128::new(v, p, s)))
.unwrap_or(Value::Null),
ScalarValue::Struct(struct_array) => {
- let struct_type = StructType::from(struct_array.fields());
+ // A struct scalar carries a single element; a null (or empty)
+ // element means the struct itself is null, for example a null
+ // struct inside a list.
+ if struct_array.is_empty() || struct_array.is_null(0) {
+ return Ok(Value::Null);
+ }
+ // Build the struct type fallibly: an unrepresentable arrow
+ // field type must surface as an error, not a panic.
+ let struct_type = StructType::try_from_arrow_fields(struct_array.fields())?;
let items = struct_array
.columns()
.iter()
@@ -1244,8 +1253,6 @@ impl TryFrom for Value {
| ScalarValue::LargeListView(_)
| ScalarValue::Union(_, _, _)
| ScalarValue::Float16(_)
- | ScalarValue::Utf8View(_)
- | ScalarValue::BinaryView(_)
| ScalarValue::Map(_)
| ScalarValue::Date64(_)
| ScalarValue::RunEndEncoded(_, _, _) => {
@@ -1259,6 +1266,21 @@ impl TryFrom for Value {
}
}
+/// Converts the value at `index` of an arrow array into a [`Value`].
+///
+/// Unlike the unchecked conversion in [`crate::vectors::StructVector`]'s `get`,
+/// this propagates conversion errors (for example arrow types greptimedb
+/// cannot represent, such as `Decimal256`) instead of panicking, so callers
+/// can surface a proper error to the user.
+pub fn try_value_from_array(array: &dyn Array, index: usize) -> Result {
+ if array.is_null(index) {
+ return Ok(Value::Null);
+ }
+ let scalar =
+ ScalarValue::try_from_array(array, index).context(ConvertArrowArrayToScalarsSnafu)?;
+ Value::try_from(scalar)
+}
+
impl From> for Value {
fn from(value: ValueRef<'_>) -> Self {
match value {
@@ -2157,6 +2179,73 @@ pub(crate) mod tests {
Value::Struct(struct_value),
scalar_struct_value.try_into().unwrap()
);
+
+ // view-typed strings and binaries convert like their non-view forms
+ assert_eq!(
+ Value::String("abc".into()),
+ ScalarValue::Utf8View(Some("abc".into()))
+ .try_into()
+ .unwrap()
+ );
+ assert_eq!(Value::Null, ScalarValue::Utf8View(None).try_into().unwrap());
+ assert_eq!(
+ Value::Binary(Bytes::from(vec![1, 2])),
+ ScalarValue::BinaryView(Some(vec![1, 2]))
+ .try_into()
+ .unwrap()
+ );
+ assert_eq!(
+ Value::Null,
+ ScalarValue::BinaryView(None).try_into().unwrap()
+ );
+
+ // a null struct scalar is a null value, not a struct of null fields
+ let null_struct = ScalarStructBuilder::new_null(build_struct_type().as_arrow_fields());
+ assert_eq!(Value::Null, null_struct.try_into().unwrap());
+ }
+
+ #[test]
+ fn test_try_value_from_array() {
+ use arrow_array::Int32Array;
+ use datafusion_common::arrow::datatypes::i256;
+
+ let array = Int32Array::from(vec![Some(1), None]);
+ assert_eq!(Value::Int32(1), try_value_from_array(&array, 0).unwrap());
+ assert_eq!(Value::Null, try_value_from_array(&array, 1).unwrap());
+
+ // a struct with representable fields converts
+ let supported = ScalarStructBuilder::new()
+ .with_name_and_scalar("i", ScalarValue::Int32(Some(7)))
+ .build()
+ .unwrap();
+ let ScalarValue::Struct(array) = supported else {
+ unreachable!();
+ };
+ assert_eq!(
+ Value::Struct(StructValue::new(
+ vec![Value::Int32(7)],
+ StructType::new(std::sync::Arc::new(vec![StructField::new(
+ "i",
+ ConcreteDataType::int32_datatype(),
+ false,
+ )]))
+ )),
+ try_value_from_array(array.as_ref(), 0).unwrap()
+ );
+
+ // an unsupported arrow field type errors instead of panicking
+ let unsupported = ScalarStructBuilder::new()
+ .with_name_and_scalar(
+ "d",
+ ScalarValue::Decimal256(Some(i256::from_i128(1)), 38, 10),
+ )
+ .build()
+ .unwrap();
+ let ScalarValue::Struct(array) = unsupported else {
+ unreachable!();
+ };
+ let error = try_value_from_array(array.as_ref(), 0).unwrap_err();
+ assert!(error.to_string().contains("Unsupported arrow data type"));
}
#[test]
diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs
index 08a2fd86555..5b6265c7b37 100644
--- a/src/datatypes/src/vectors/json/array.rs
+++ b/src/datatypes/src/vectors/json/array.rs
@@ -26,15 +26,13 @@ use serde_json::Value;
use snafu::{OptionExt, ResultExt};
use crate::arrow_array::{binary_array_value, string_array_value};
-use crate::data_type::ConcreteDataType;
+use crate::data_type::{ConcreteDataType, DataType as _};
use crate::error::{
AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result,
};
use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field};
-use crate::json::JsonSettings;
-use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb};
-use crate::prelude::{DataType as _, Value as GreptimeValue};
-use crate::value::{ListValue, StructValue};
+use crate::json::value::decode_json_variant;
+use crate::json::{JsonSettings, TypeHintMismatchPolicy, coerce_json_value_to_type};
use crate::vectors::MutableVector;
use crate::vectors::json::builder::{JsonVectorBuilder, json2_physical_data_type};
use crate::vectors::json::variant::variant_to_json_values;
@@ -123,6 +121,23 @@ impl JsonArray<'_> {
field: &Field,
logical_settings: &JsonSettings,
target_layout: &JsonSettings,
+ ) -> Result {
+ self.rewrite_to_v2_with_type_hint_mismatch_policy(
+ field,
+ logical_settings,
+ target_layout,
+ TypeHintMismatchPolicy::Reject,
+ )
+ }
+
+ /// Rewrites a JSON2 array to the specified v2 physical layout using the
+ /// given type hint mismatch policy.
+ pub fn rewrite_to_v2_with_type_hint_mismatch_policy(
+ &self,
+ field: &Field,
+ logical_settings: &JsonSettings,
+ target_layout: &JsonSettings,
+ policy: TypeHintMismatchPolicy,
) -> Result {
let is_v2 = json2_remainder_field(field)?.is_some();
if is_v2 && self.inner.data_type() == &json2_physical_data_type(target_layout) {
@@ -141,7 +156,8 @@ impl JsonArray<'_> {
if value.is_null() {
builder.push_null();
} else {
- let value = logical_settings.encode(value)?;
+ let value =
+ logical_settings.encode_with_type_hint_mismatch_policy(value, policy)?;
builder.try_push_value_ref(&value.as_value_ref())?;
}
}
@@ -319,87 +335,12 @@ fn project_json_values(values: Vec, to_type: &DataType) -> Result Result {
- if value.is_null() {
- return Ok(GreptimeValue::Null);
- }
-
- if to_type.is_string() {
- let value = match value {
- Value::String(value) => value,
- value => value.to_string(),
- };
- return Ok(GreptimeValue::String(value.into()));
- }
-
- if matches!(to_type, ConcreteDataType::Binary(_)) {
- return Ok(GreptimeValue::Binary(
- encode_serde_json_as_jsonb(value).into(),
- ));
- }
-
- if let Some(struct_type) = to_type.as_struct() {
- let Value::Object(mut object) = value else {
- return Ok(GreptimeValue::Null);
- };
- let values = struct_type
- .fields()
- .iter()
- .map(|field| {
- object
- .remove(field.name())
- .map(|value| project_json_value_to_type(value, field.data_type()))
- .transpose()
- .map(|value| value.unwrap_or(GreptimeValue::Null))
- })
- .collect::>>()?;
- return Ok(GreptimeValue::Struct(StructValue::new(
- values,
- struct_type.clone(),
- )));
- }
-
- if let Some(list_type) = to_type.as_list() {
- let Value::Array(values) = value else {
- return Ok(GreptimeValue::Null);
- };
- let item_type = list_type.item_type().clone();
- let values = values
- .into_iter()
- .map(|value| project_json_value_to_type(value, &item_type))
- .collect::>>()?;
- return Ok(GreptimeValue::List(ListValue::new(
- values,
- Arc::new(item_type),
- )));
- }
-
- let value = match value {
- Value::Bool(value) => GreptimeValue::Boolean(value),
- Value::Number(value) => {
- if let Some(value) = value.as_i64() {
- GreptimeValue::Int64(value)
- } else if let Some(value) = value.as_u64() {
- GreptimeValue::UInt64(value)
- } else if let Some(value) = value.as_f64() {
- GreptimeValue::Float64(value.into())
- } else {
- GreptimeValue::Null
- }
- }
- Value::String(value) => GreptimeValue::String(value.into()),
- Value::Array(_) | Value::Object(_) => GreptimeValue::Null,
- Value::Null => GreptimeValue::Null,
- };
- Ok(to_type.try_cast(value).unwrap_or(GreptimeValue::Null))
-}
-
impl<'a> From<&'a ArrayRef> for JsonArray<'a> {
fn from(inner: &'a ArrayRef) -> Self {
Self { inner }
diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md
index d552d73cf39..49a4b8fcd0f 100644
--- a/src/frontend/AGENTS.md
+++ b/src/frontend/AGENTS.md
@@ -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
diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs
index 86ba2cb9bfb..72bf322b1f1 100644
--- a/src/frontend/src/frontend.rs
+++ b/src/frontend/src/frontend.rs
@@ -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 =
Pin> + Send + Sync + 'static>>;
+ #[test]
+ fn test_logical_batcher_protocols_from_env() {
+ temp_env::with_vars(
+ [
+ (
+ "FRONTEND_LOGICAL_TEST__PENDING_ROWS_BATCHER__PROTOCOLS",
+ Some("otlp,influxdb"),
+ ),
+ (
+ "FRONTEND_LOGICAL_TEST__PENDING_ROWS_BATCHER__LOGICAL_TABLE__PROTOCOLS",
+ Some("prom,otlp"),
+ ),
+ ],
+ || {
+ let options =
+ FrontendOptions::load_layered_options(None, "FRONTEND_LOGICAL_TEST").unwrap();
+ assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2);
+ assert_eq!(
+ options
+ .pending_rows_batcher
+ .logical_table
+ .unwrap()
+ .protocols,
+ vec![
+ 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);
diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs
index 22d2afaa7bc..92bc74ca10b 100644
--- a/src/frontend/src/instance.rs
+++ b/src/frontend/src/instance.rs
@@ -21,6 +21,7 @@ mod import_packed;
mod influxdb;
mod jaeger;
mod log_handler;
+mod logical_batcher;
mod logs;
mod opentsdb;
mod otlp;
@@ -31,7 +32,7 @@ mod region_query;
use std::collections::HashSet;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
-use std::sync::{Arc, atomic};
+use std::sync::{Arc, OnceLock, atomic};
use std::time::{Duration, SystemTime};
use async_stream::stream;
@@ -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>>>,
deleter: DeleterRef,
table_metadata_manager: TableMetadataManagerRef,
event_recorder: EventRecorderRef,
diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs
index 01e9ef034e3..d9b86ac5a05 100644
--- a/src/frontend/src/instance/builder.rs
+++ b/src/frontend/src/instance/builder.rs
@@ -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> {
- if !options.pending_rows_batching_enabled()
- || (self.options.prom_store.with_metric_engine
- && options
- .protocols
- .iter()
- .all(|protocol| *protocol == BatchingProtocol::Prom))
- {
- return None;
- }
- TablePendingRowsBatcher::try_new(
- options.pending_rows_flush_interval,
- options.max_batch_rows,
- options.max_concurrent_flushes,
- options.worker_channel_capacity,
- options.max_inflight_requests,
- options.flow_notification_queue_capacity,
- bulk_inserter.clone(),
- )
- .map(|batcher| batcher as Arc)
- };
+ let build_batcher = |options: &BatcherOptions| -> Option> {
+ if !options.pending_rows_batching_enabled()
+ || (self.options.prom_store.with_metric_engine
+ && options
+ .protocols
+ .iter()
+ .all(|protocol| *protocol == BatchingProtocol::Prom))
+ {
+ return None;
+ }
+ TablePendingRowsBatcher::try_new(
+ options.pending_rows_flush_interval,
+ options.max_batch_rows,
+ options.max_concurrent_flushes,
+ options.worker_channel_capacity,
+ options.max_inflight_requests,
+ options.flow_notification_queue_capacity,
+ bulk_inserter.clone(),
+ )
+ .map(|batcher| batcher as Arc)
+ };
let inserter = Arc::new(
create_inserter()
- .with_pending_rows_batcher(build_batcher(&self.options.pending_rows_batcher)),
+ .with_pending_rows_batcher(build_batcher(self.options.table_batcher_options())),
);
let deleter = Arc::new(Deleter::new(
self.catalog_manager.clone(),
@@ -382,10 +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,
diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs
new file mode 100644
index 00000000000..28ed19b3e14
--- /dev/null
+++ b/src/frontend/src/instance/logical_batcher.rs
@@ -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, 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> {
+ self.logical_batcher.get().and_then(Option::as_ref)
+ }
+}
+
+// The instance owns the shared batcher; schema preparation must not keep that
+// owner alive through a reference cycle.
+struct LogicalTables(Weak);
+
+#[async_trait]
+impl PendingRowsSchemaAlterer for LogicalTables {
+ async fn create_tables_if_missing_batch(
+ &self,
+ catalog: &str,
+ schema: &str,
+ tables: &[(&str, &[ColumnSchema])],
+ with_metric_engine: bool,
+ ctx: QueryContextRef,
+ ) -> Result<()> {
+ self.0
+ .upgrade()
+ .context(BatcherChannelClosedSnafu)?
+ .create_tables_if_missing_batch(catalog, schema, tables, with_metric_engine, ctx)
+ .await
+ }
+
+ async fn add_missing_prom_tag_columns_batch(
+ &self,
+ catalog: &str,
+ schema: &str,
+ tables: &[(&str, &[String])],
+ ctx: QueryContextRef,
+ ) -> Result<()> {
+ self.0
+ .upgrade()
+ .context(BatcherChannelClosedSnafu)?
+ .add_missing_prom_tag_columns_batch(catalog, schema, tables, ctx)
+ .await
+ }
+}
diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs
index fa40c7eb59a..7a0f706c58d 100644
--- a/src/frontend/src/instance/otlp.rs
+++ b/src/frontend/src/instance/otlp.rs
@@ -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,
diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs
index 017bba53fba..0a6612ba31c 100644
--- a/src/frontend/src/server.rs
+++ b/src/frontend/src/server.rs
@@ -24,9 +24,7 @@ use common_base::Plugins;
use common_config::Configurable;
use common_telemetry::{info, warn};
use meta_client::MetaClientOptions;
-use servers::batcher::logical_table::{
- LogicalTablePendingRowsBatcher, pending_rows_batch_sync_enabled,
-};
+use servers::batcher::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, 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),
diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs
index bcd1097cfcf..2c761c4242a 100644
--- a/src/frontend/src/service_config.rs
+++ b/src/frontend/src/service_config.rs
@@ -26,6 +26,6 @@ pub use jaeger::JaegerOptions;
pub use mysql::MysqlOptions;
pub use opentsdb::OpentsdbOptions;
pub use otlp::OtlpOptions;
-pub use pending_rows_batcher::PendingRowsBatcherOptions;
+pub use pending_rows_batcher::{BatcherOptions, PendingRowsBatcherOptions};
pub use postgres::PostgresOptions;
pub use prom_store::PromStoreOptions;
diff --git a/src/frontend/src/service_config/pending_rows_batcher.rs b/src/frontend/src/service_config/pending_rows_batcher.rs
index ef7ce14ef02..61572015d7c 100644
--- a/src/frontend/src/service_config/pending_rows_batcher.rs
+++ b/src/frontend/src/service_config/pending_rows_batcher.rs
@@ -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,
+}
+
+/// 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,
/// 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, D::Error>
+where
+ D: serde::Deserializer<'de>,
+{
+ let options = Option::::deserialize(deserializer)?;
+ if options.as_ref().is_some_and(|options| {
+ options
+ .protocols
+ .iter()
+ .any(|protocol| !matches!(protocol, BatchingProtocol::Prom | BatchingProtocol::Otlp))
+ }) {
+ return Err(serde::de::Error::custom(
+ "pending_rows_batcher.logical_table only supports prom and otlp",
+ ));
+ }
+ Ok(options)
+}
+
+impl FrontendOptions {
+ /// Explicit configuration wins even when it disables batching.
+ pub(crate) fn table_batcher_options(&self) -> &BatcherOptions {
+ &self.pending_rows_batcher.table
+ }
+
+ /// Resolve legacy Prom controls once without merging fields across config blocks.
+ pub(crate) fn logical_batcher_options(&self) -> BatcherOptions {
+ if let Some(options) = &self.pending_rows_batcher.logical_table {
+ return options.clone();
+ }
+ let shared = &self.pending_rows_batcher.table;
+ if shared.protocols.contains(&BatchingProtocol::Prom)
+ && shared.pending_rows_batching_enabled()
+ {
+ let mut options = shared.clone();
+ options.protocols = vec![BatchingProtocol::Prom];
+ return options;
+ }
+ let prom = &self.prom_store;
+ BatcherOptions {
+ protocols: vec![BatchingProtocol::Prom],
+ pending_rows_flush_interval: prom.pending_rows_flush_interval,
+ max_batch_rows: prom.max_batch_rows,
+ max_concurrent_flushes: prom.max_concurrent_flushes,
+ worker_channel_capacity: prom.worker_channel_capacity,
+ max_inflight_requests: prom.max_inflight_requests,
+ flow_notification_queue_capacity: prom.flow_notification_queue_capacity,
+ }
+ }
+}
+
#[cfg(test)]
mod tests {
use crate::service_config::pending_rows_batcher::*;
+ #[test]
+ fn test_batcher_config_precedence() {
+ use crate::frontend::FrontendOptions;
+ let legacy = "[prom_store]\nenable = true\nwith_metric_engine = true\npending_rows_flush_interval = '2s'\nmax_batch_rows = 9\n";
+ let shared = "[pending_rows_batcher]\nprotocols = ['prom', 'otlp']\npending_rows_flush_interval = '3s'\n";
+ let options: FrontendOptions = toml::from_str(legacy).unwrap();
+ assert_eq!(options.logical_batcher_options().max_batch_rows, 9);
+ assert_eq!(
+ options.logical_batcher_options().protocols,
+ vec![BatchingProtocol::Prom]
+ );
+ let options: FrontendOptions = toml::from_str(&format!("{legacy}{shared}")).unwrap();
+ assert_eq!(
+ options
+ .logical_batcher_options()
+ .pending_rows_flush_interval,
+ Duration::from_secs(3)
+ );
+ assert_eq!(options.logical_batcher_options().max_batch_rows, 100_000);
+ assert_eq!(
+ options.logical_batcher_options().protocols,
+ vec![BatchingProtocol::Prom]
+ );
+ for explicit in [
+ "",
+ "protocols = ['otlp']\npending_rows_flush_interval = '5s'",
+ ] {
+ let options: FrontendOptions = toml::from_str(&format!(
+ "{legacy}{shared}[pending_rows_batcher.logical_table]\n{explicit}"
+ ))
+ .unwrap();
+ assert_eq!(
+ &options.logical_batcher_options(),
+ options.pending_rows_batcher.logical_table.as_ref().unwrap()
+ );
+ assert_eq!(
+ options.table_batcher_options(),
+ &toml::from_str::(&format!("{legacy}{shared}"))
+ .unwrap()
+ .pending_rows_batcher
+ .table
+ );
+ let restored: FrontendOptions =
+ toml::from_str(&toml::to_string(&options).unwrap()).unwrap();
+ assert_eq!(
+ options.pending_rows_batcher.logical_table,
+ restored.pending_rows_batcher.logical_table
+ );
+ assert_eq!(options.pending_rows_batcher, restored.pending_rows_batcher);
+ }
+ }
+
+ #[test]
+ fn test_batcher_config_independent_defaults() {
+ use crate::frontend::FrontendOptions;
+
+ let parent = "[pending_rows_batcher]\nprotocols = ['prom']\npending_rows_flush_interval = '3s'\nmax_batch_rows = 7\n";
+ let options: FrontendOptions = toml::from_str(parent).unwrap();
+ let serialized = toml::to_string(&options).unwrap();
+ assert!(!serialized.contains("[pending_rows_batcher.logical_table]"));
+ assert!(!serialized.contains("[pending_rows_batcher.table]"));
+ let restored: FrontendOptions = toml::from_str(&serialized).unwrap();
+ assert_eq!(restored.logical_batcher_options().max_batch_rows, 7);
+
+ let options: FrontendOptions =
+ toml::from_str(&format!("{parent}[pending_rows_batcher.logical_table]")).unwrap();
+ assert_eq!(options.table_batcher_options().max_batch_rows, 7);
+ assert_eq!(options.logical_batcher_options(), BatcherOptions::default());
+ assert!(
+ !options
+ .logical_batcher_options()
+ .pending_rows_batching_enabled()
+ );
+ }
+
+ #[test]
+ fn test_logical_batcher_config_protocols() {
+ use crate::frontend::FrontendOptions;
+ for protocols in ["[]", "['prom']", "['otlp']", "['prom', 'otlp']"] {
+ let options: FrontendOptions = toml::from_str(&format!(
+ "[pending_rows_batcher.logical_table]\nprotocols = {protocols}"
+ ))
+ .unwrap();
+ assert!(options.pending_rows_batcher.logical_table.is_some());
+ }
+ for protocol in [
+ "influxdb",
+ "logs",
+ "loki",
+ "http_sql",
+ "opentsdb",
+ "elasticsearch",
+ "splunk",
+ ] {
+ assert!(
+ toml::from_str::(&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::(&format!("protocols = ['{invalid}']"))
- .is_err()
+ toml::from_str::(&format!("protocols = ['{invalid}']")).is_err()
);
}
}
#[test]
fn test_notification_capacity() {
- let default = PendingRowsBatcherOptions::default();
+ let default = BatcherOptions::default();
assert_eq!(default.flow_notification_queue_capacity.get(), 1024);
- let configured: PendingRowsBatcherOptions =
+ let configured: BatcherOptions =
toml::from_str("flow_notification_queue_capacity = 8").unwrap();
assert_eq!(configured.flow_notification_queue_capacity.get(), 8);
- assert!(
- toml::from_str::("flow_notification_queue_capacity = 0")
- .is_err()
- );
+ assert!(toml::from_str::("flow_notification_queue_capacity = 0").is_err());
}
#[test]
fn test_defaults_and_roundtrip() {
- let options: PendingRowsBatcherOptions = toml::from_str("").unwrap();
- assert_eq!(options, PendingRowsBatcherOptions::default());
+ let options: BatcherOptions = toml::from_str("").unwrap();
+ assert_eq!(options, BatcherOptions::default());
assert!(!options.pending_rows_batching_enabled());
assert_eq!(options.max_batch_rows, 100_000);
assert_eq!(options.max_concurrent_flushes, 256);
- assert_eq!(options.worker_channel_capacity, 65526);
+ assert_eq!(options.worker_channel_capacity, 65_536);
assert_eq!(options.max_inflight_requests, 3000);
let serialized = toml::to_string(&options).unwrap();
assert_eq!(
options,
- toml::from_str::(&serialized).unwrap()
+ toml::from_str::(&serialized).unwrap()
);
}
#[test]
fn test_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();
diff --git a/src/frontend/src/service_config/prom_store.rs b/src/frontend/src/service_config/prom_store.rs
index f8429a86925..2fdd05043db 100644
--- a/src/frontend/src/service_config/prom_store.rs
+++ b/src/frontend/src/service_config/prom_store.rs
@@ -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 {
diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs
index b40f06e3cfe..cbc859716a4 100644
--- a/src/meta-srv/src/metasrv/builder.rs
+++ b/src/meta-srv/src/metasrv/builder.rs
@@ -250,6 +250,7 @@ impl MetasrvBuilder {
let event_recorder = Arc::new(EventRecorderImpl::with_event_type_filter(
Box::new(EventHandlerImpl::new(event_inserter)),
options.event_recorder.event_types.clone(),
+ options.event_recorder.flush_interval,
));
let selector = selector.unwrap_or_else(|| Arc::new(LeaseBasedSelector));
diff --git a/src/metric-engine/src/data_region.rs b/src/metric-engine/src/data_region.rs
index 55b9eaaf683..a11b3e57b54 100644
--- a/src/metric-engine/src/data_region.rs
+++ b/src/metric-engine/src/data_region.rs
@@ -227,6 +227,7 @@ impl DataRegion {
| AlterKind::UnsetRegionOptions { keys: _ }
| AlterKind::SetIndexes { options: _ }
| AlterKind::UnsetIndexes { options: _ }
+ | AlterKind::SetJsonSettings { .. }
| AlterKind::SyncColumns {
column_metadatas: _,
} => {
diff --git a/src/mito-codec/src/row_converter.rs b/src/mito-codec/src/row_converter.rs
index fae2997182c..60bbf1935e0 100644
--- a/src/mito-codec/src/row_converter.rs
+++ b/src/mito-codec/src/row_converter.rs
@@ -134,6 +134,11 @@ pub trait PrimaryKeyCodec: Send + Sync + Debug {
/// Returns the encoding type of the primary key.
fn encoding(&self) -> PrimaryKeyEncoding;
+ /// Returns the dense codec for schema-aware positional access, if applicable.
+ fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> {
+ None
+ }
+
/// Decodes the primary key from the given bytes.
///
/// Returns a [`CompositeValues`] that follows the primary key ordering.
@@ -166,3 +171,28 @@ pub fn build_primary_key_codec_with_fields(
}
}
}
+
+/// Finds the checked boundary of an Option, shared by Dense and Sparse.
+/// This validates framing, not UTF-8; consumers decoding strings validate UTF-8.
+pub(crate) fn encoded_string_len(bytes: &[u8]) -> memcomparable::Result {
+ match bytes.first().copied().ok_or(memcomparable::Error::Eof)? {
+ 0 => return Ok(1),
+ 1 => {}
+ marker => return Err(memcomparable::Error::InvalidTagEncoding(marker as usize)),
+ }
+ match bytes.get(1).copied().ok_or(memcomparable::Error::Eof)? {
+ 0 => return Ok(2),
+ 1 => {}
+ marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)),
+ }
+ let mut end = 2;
+ loop {
+ let chunk = bytes.get(end..end + 9).ok_or(memcomparable::Error::Eof)?;
+ end += 9;
+ match chunk[8] {
+ 1..=8 => return Ok(end),
+ 9 => {}
+ marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)),
+ }
+ }
+}
diff --git a/src/mito-codec/src/row_converter/dense.rs b/src/mito-codec/src/row_converter/dense.rs
index 56289a08b11..ad85f4ac8f1 100644
--- a/src/mito-codec/src/row_converter/dense.rs
+++ b/src/mito-codec/src/row_converter/dense.rs
@@ -27,7 +27,7 @@ use datatypes::value::ValueRef;
use memcomparable::{Deserializer, Serializer};
use paste::paste;
use serde::{Deserialize, Serialize};
-use snafu::ResultExt;
+use snafu::{ResultExt, ensure};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::{RegionMetadata, RegionMetadataRef};
use store_api::storage::ColumnId;
@@ -38,7 +38,7 @@ use crate::error::{
use crate::key_values::KeyValue;
use crate::primary_key_filter::DensePrimaryKeyFilter;
use crate::row_converter::{
- CompositeValues, PrimaryKeyCodec, PrimaryKeyCodecExt, PrimaryKeyFilter,
+ CompositeValues, PrimaryKeyCodec, PrimaryKeyCodecExt, PrimaryKeyFilter, encoded_string_len,
};
/// Field to serialize and deserialize value in memcomparable format.
@@ -290,84 +290,112 @@ impl SortField {
bytes: &[u8],
deserializer: &mut Deserializer<&[u8]>,
) -> Result {
- let len = self.encoded_field_len(&bytes[deserializer.position()..])?;
+ let pos = deserializer.position();
+ let remaining = bytes
+ .get(pos..)
+ .ok_or(memcomparable::Error::Eof)
+ .context(error::DeserializeFieldSnafu)?;
+ let len = Self::encoded_len(self.encode_data_type(), remaining)?;
+ if len > 1 && self.encode_data_type().is_boolean() && remaining[1] > 1 {
+ return Err(memcomparable::Error::InvalidBoolEncoding(remaining[1]))
+ .context(error::DeserializeFieldSnafu);
+ }
deserializer.advance(len);
Ok(len)
}
/// Checks field boundaries before any unchecked reads in memcomparable.
- fn encoded_field_len(&self, bytes: &[u8]) -> Result {
- let invalid = || {
- error::InvalidDensePrimaryKeySnafu {
- reason: "truncated field or invalid encoding",
+ fn encoded_len(data_type: &ConcreteDataType, bytes: &[u8]) -> Result {
+ let marker = bytes
+ .first()
+ .copied()
+ .ok_or(memcomparable::Error::Eof)
+ .context(error::DeserializeFieldSnafu)?;
+ match marker {
+ 0 => return Ok(1),
+ 1 => {}
+ value => {
+ return Err(memcomparable::Error::InvalidTagEncoding(value as usize))
+ .context(error::DeserializeFieldSnafu);
}
- .build()
- };
- match bytes.first() {
- Some(0) => return Ok(1),
- Some(1) => {}
- _ => return Err(invalid()),
}
- let len = match self.encode_data_type() {
- ConcreteDataType::Boolean(_)
- | ConcreteDataType::Int8(_)
- | ConcreteDataType::UInt8(_) => 2,
+ let to_skip = match data_type {
+ ConcreteDataType::Boolean(_) => 2,
+ ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => 2,
ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => 3,
- ConcreteDataType::Int32(_)
- | ConcreteDataType::UInt32(_)
- | ConcreteDataType::Float32(_)
- | ConcreteDataType::Date(_) => 5,
- ConcreteDataType::Int64(_)
- | ConcreteDataType::UInt64(_)
- | ConcreteDataType::Float64(_)
- | ConcreteDataType::Timestamp(_) => 9,
- ConcreteDataType::Time(_) | ConcreteDataType::Duration(_) => 10,
+ ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => 5,
+ ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => 9,
+ ConcreteDataType::Float32(_) => 5,
+ ConcreteDataType::Float64(_) => 9,
+ ConcreteDataType::Binary(_)
+ | ConcreteDataType::Json(_)
+ | ConcreteDataType::Vector(_) => {
+ // Binary is encoded as a sequence of bytes, not chunked strings.
+ return encoded_binary_len(bytes).context(error::DeserializeFieldSnafu);
+ }
+ ConcreteDataType::String(_) => {
+ return encoded_string_len(bytes).context(error::DeserializeFieldSnafu);
+ }
+ ConcreteDataType::Date(_) => 5,
+ ConcreteDataType::Timestamp(_) => 9, // We treat timestamp as Option
+ ConcreteDataType::Time(_) => 10, // i64 and 1 byte time unit
+ ConcreteDataType::Duration(_) => 10,
ConcreteDataType::Interval(IntervalType::YearMonth(_)) => 5,
ConcreteDataType::Interval(IntervalType::DayTime(_)) => 9,
ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => 17,
ConcreteDataType::Decimal128(_) => 19,
- ConcreteDataType::Binary(_)
- | ConcreteDataType::Json(_)
- | ConcreteDataType::Vector(_) => {
- // A serde sequence: each byte is preceded by 1, terminated by 0.
- let mut pos = 1;
- while bytes.get(pos) == Some(&1) {
- pos += 2;
- }
- if bytes.get(pos) != Some(&0) {
- return Err(invalid());
- }
- pos + 1
- }
- ConcreteDataType::String(_) => {
- match bytes.get(1) {
- Some(0) => return Ok(2),
- Some(1) => {}
- _ => return Err(invalid()),
- }
- let mut pos = 2;
- loop {
- match bytes.get(pos + 8) {
- Some(1..=8) => break pos + 9,
- Some(9) => pos += 9,
- _ => return Err(invalid()),
- }
- }
- }
- data_type => {
- return error::NotSupportedFieldSnafu {
+ ConcreteDataType::Null(_)
+ | ConcreteDataType::List(_)
+ | ConcreteDataType::Struct(_)
+ | ConcreteDataType::Dictionary(_) => {
+ return NotSupportedFieldSnafu {
data_type: data_type.clone(),
}
.fail();
}
};
- snafu::ensure!(
- bytes.len() >= len,
- error::InvalidDensePrimaryKeySnafu {
- reason: "truncated field",
- }
- );
- Ok(len)
+ if bytes.len() < to_skip {
+ return Err(memcomparable::Error::Eof).context(error::DeserializeFieldSnafu);
+ }
+ Ok(to_skip)
+ }
+
+ /// Decodes primitive fields straight into a reference value for column
+ /// builders, avoiding owned Value construction and conversion. The caller has
+ /// already checked the encoded boundary.
+ fn deserialize_primitive_ref(
+ data_type: &ConcreteDataType,
+ deserializer: &mut Deserializer,
+ ) -> Option>> {
+ macro_rules! decode_primitive {
+ ($($variant:ident, $native:ty),* $(,)?) => {
+ match data_type {
+ $(ConcreteDataType::$variant(_) => Some(
+ Option::<$native>::deserialize(deserializer)
+ .map(ValueRef::from)
+ .context(error::DeserializeFieldSnafu)
+ ),)*
+ _ => None,
+ }
+ };
+ }
+ decode_primitive!(
+ Boolean, bool, Int8, i8, Int16, i16, Int32, i32, Int64, i64, UInt8, u8, UInt16, u16,
+ UInt32, u32, UInt64, u64, Float32, f32, Float64, f64,
+ )
+ }
+}
+
+/// The Option marker has already been checked by SortField.
+fn encoded_binary_len(bytes: &[u8]) -> memcomparable::Result {
+ let mut current = 1;
+ loop {
+ match bytes.get(current).copied() {
+ Some(0) => return Ok(current + 1),
+ Some(1) if bytes.get(current + 1).is_some() => current += 2,
+ Some(1) | None => return Err(memcomparable::Error::Eof),
+ Some(marker) => return Err(memcomparable::Error::InvalidSeqEncoding(marker)),
+ }
}
}
@@ -444,7 +472,16 @@ impl DensePrimaryKeyCodec {
return Ok(index);
}
let start = deserializer.position();
- let len = field.encoded_field_len(&bytes[start..])?;
+ // Preserve the prefix error contract used by primary-key range mapping.
+ let len = SortField::encoded_len(field.encode_data_type(), &bytes[start..]).map_err(
+ |source| match source {
+ error::Error::DeserializeField { .. } => error::InvalidDensePrimaryKeySnafu {
+ reason: "truncated field or invalid encoding",
+ }
+ .build(),
+ source => source,
+ },
+ )?;
// Production range mapping only needs boundaries, not allocated field values.
#[cfg(any(debug_assertions, test))]
field.deserialize(&mut Deserializer::new(&bytes[start..start + len]))?;
@@ -470,6 +507,73 @@ impl DensePrimaryKeyCodec {
Ok(values)
}
+ /// Iterates over all source PK fields in schema order, checking each field's
+ /// encoded boundary before deserializing it. Stops after the first error.
+ ///
+ /// Unlike positional access, this requires no offsets or value cache. Callers
+ /// can consume each value immediately instead of retaining the entire key.
+ pub fn decode_dense_iter<'a>(
+ &'a self,
+ bytes: &'a [u8],
+ ) -> impl Iterator- > + 'a {
+ self.ordered_primary_key_columns.iter().scan(
+ Some(Deserializer::new(bytes)),
+ move |state, (id, field)| {
+ let deserializer = state.as_mut()?;
+ let remaining = &bytes[deserializer.position()..];
+ let decoded = SortField::encoded_len(field.encode_data_type(), remaining)
+ .and_then(|_| field.deserialize(deserializer).map(|value| (*id, value)));
+ if decoded.is_err() {
+ *state = None;
+ }
+ Some(decoded)
+ },
+ )
+ }
+
+ /// Returns the column ids and sort fields in encoded primary-key order.
+ pub fn fields(&self) -> &[(ColumnId, SortField)] {
+ &self.ordered_primary_key_columns
+ }
+
+ /// Decodes all fields in source order into a consumer, checking boundaries.
+ /// Strings borrow the reusable buffer for the duration of each callback, so
+ /// materializing them into column builders needs no per-value owned string.
+ pub fn decode_dense_with(
+ &self,
+ bytes: &[u8],
+ value_buf: &mut Vec,
+ mut consume: impl FnMut(usize, ValueRef<'_>),
+ ) -> Result<()> {
+ let mut deserializer = Deserializer::new(bytes);
+ for (pos, (_, field)) in self.ordered_primary_key_columns.iter().enumerate() {
+ let data_type = field.encode_data_type();
+ let remaining = &bytes[deserializer.position()..];
+ SortField::encoded_len(data_type, remaining)?;
+ if data_type.is_string() && remaining[0] != 0 {
+ deserializer.advance(1);
+ deserializer
+ .read_bytes_into(value_buf)
+ .context(error::DeserializeFieldSnafu)?;
+ let value = std::str::from_utf8(value_buf).map_err(|err| {
+ error::InvalidDensePrimaryKeySnafu {
+ reason: format!("string is not valid UTF-8: {err}"),
+ }
+ .build()
+ })?;
+ consume(pos, ValueRef::String(value));
+ } else if let Some(value) =
+ SortField::deserialize_primitive_ref(data_type, &mut deserializer)
+ {
+ consume(pos, value?);
+ } else {
+ let value = field.deserialize(&mut deserializer)?;
+ consume(pos, value.as_value_ref());
+ }
+ }
+ Ok(())
+ }
+
/// Returns the field at `pos`.
///
/// # Panics
@@ -488,61 +592,61 @@ impl DensePrimaryKeyCodec {
offsets_buf: &mut Vec,
deserializer: &mut Deserializer<&[u8]>,
) -> Result {
- if pos < offsets_buf.len() {
- // We computed the offset before.
- let offset = offsets_buf[pos];
- deserializer.advance(offset);
- return Ok(offset);
- }
-
+ ensure!(
+ pos < self.num_fields(),
+ error::InvalidDensePrimaryKeySnafu {
+ reason: format!(
+ "field position {pos} exceeds field count {}",
+ self.num_fields()
+ ),
+ }
+ );
if offsets_buf.is_empty() {
- let mut offset = 0;
- // Skip values before `pos`.
- for i in 0..pos {
- // Offset to skip before reading value i.
- offsets_buf.push(offset);
- let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?;
- offset += skip;
- }
- // Offset to skip before reading this value.
- offsets_buf.push(offset);
- Ok(offset)
- } else {
- // Offsets are not enough.
- let value_start = offsets_buf.len() - 1;
- // Advances to decode value at `value_start`.
- let mut offset = offsets_buf[value_start];
- deserializer.advance(offset);
- for i in value_start..pos {
- // Skip value i.
- let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?;
- // Offset for the value at i + 1.
- offset += skip;
- offsets_buf.push(offset);
- }
- Ok(offset)
+ // The schema bounds this cache. Reserve once even when the first
+ // requested field is near the end of a wide key.
+ offsets_buf.reserve_exact(self.num_fields());
+ offsets_buf.push(0);
}
+ // Start at the requested field if cached, otherwise resume discovery
+ // from the furthest known boundary.
+ let value_start = pos.min(offsets_buf.len() - 1);
+ let mut offset = offsets_buf[value_start];
+ ensure!(
+ offset <= bytes.len(),
+ error::InvalidDensePrimaryKeySnafu {
+ reason: "cached offset exceeds key length",
+ }
+ );
+ deserializer.advance(offset);
+ for i in value_start..pos {
+ offset += self.field_at(i).skip_deserialize(bytes, deserializer)?;
+ offsets_buf.push(offset);
+ }
+ Ok(offset)
}
/// Decode value at `pos` in `bytes`.
///
- /// The i-th element in offsets buffer is how many bytes to skip in order to read value at `pos`.
+ /// The i-th element in the offsets buffer is the start of field i. The buffer
+ /// must only be reused for the same encoded key and codec.
+ #[inline]
pub fn decode_value_at(
&self,
bytes: &[u8],
pos: usize,
offsets_buf: &mut Vec,
) -> Result {
- let mut deserializer = Deserializer::new(bytes);
- self.advance_to_value_at(bytes, pos, offsets_buf, &mut deserializer)?;
-
- self.field_at(pos).deserialize(&mut deserializer)
+ let encoded = self.encoded_value_at(bytes, pos, offsets_buf)?;
+ self.field_at(pos)
+ .deserialize(&mut Deserializer::new(encoded))
}
/// Returns the encoded bytes at `pos` in `bytes`.
///
- /// The i-th element in offsets buffer is how many bytes to skip in order to read value at
- /// `pos`.
+ /// The i-th element in the offsets buffer is the start of field i. The buffer
+ /// must only be reused for the same encoded key and codec. Validates the
+ /// framing up to this field, without validating or decoding later fields.
+ #[inline]
pub fn encoded_value_at<'a>(
&self,
bytes: &'a [u8],
@@ -555,6 +659,11 @@ impl DensePrimaryKeyCodec {
let len = self
.field_at(pos)
.skip_deserialize(bytes, &mut deserializer)?;
+ // We already found the next field's start while validating this one.
+ // Reuse it instead of scanning this field again on the next lookup.
+ if offsets_buf.len() == pos + 1 && pos + 1 < self.num_fields() {
+ offsets_buf.push(offset + len);
+ }
Ok(&bytes[offset..offset + len])
}
@@ -600,6 +709,10 @@ impl PrimaryKeyCodec for DensePrimaryKeyCodec {
PrimaryKeyEncoding::Dense
}
+ fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> {
+ Some(self)
+ }
+
fn primary_key_filter(
&self,
metadata: &RegionMetadataRef,
@@ -663,6 +776,78 @@ mod tests {
}
let decoded = encoder.decode(&result).unwrap().into_dense();
assert_eq!(decoded, row);
+ let mut value_buf = Vec::new();
+ let mut borrowed = Vec::new();
+ encoder
+ .decode_dense_with(&result, &mut value_buf, |_, value| {
+ borrowed.push(Value::from(value))
+ })
+ .unwrap();
+ assert_eq!(borrowed, row);
+ let sequential = encoder
+ .decode_dense_iter(&result)
+ .collect::>>()
+ .unwrap();
+ assert_eq!(
+ sequential
+ .iter()
+ .map(|(_, value)| value)
+ .collect::>(),
+ row.iter().collect::>()
+ );
+ for end in 0..result.len() {
+ assert!(
+ encoder
+ .decode_dense_with(&result[..end], &mut value_buf, |_, _| {})
+ .is_err()
+ );
+ assert!(
+ encoder
+ .decode_dense_iter(&result[..end])
+ .collect::>>()
+ .is_err(),
+ "truncated at {end}"
+ );
+ }
+ // Every supported type must have the same boundaries in positional and
+ // full decoding. Exercise every truncation both at the requested field
+ // and while skipping it to reach a later field.
+ let mut offsets = Vec::new();
+ let mut end = 0;
+ for (pos, value) in row.iter().enumerate() {
+ let encoded = encoder
+ .encoded_value_at(&result, pos, &mut offsets)
+ .unwrap();
+ assert_eq!(
+ encoder
+ .field_at(pos)
+ .deserialize(&mut Deserializer::new(encoded))
+ .unwrap(),
+ *value
+ );
+ let start = end;
+ end += encoded.len();
+ for len in start..end {
+ assert!(
+ encoder
+ .encoded_value_at(&result[..len], pos, &mut Vec::new())
+ .is_err()
+ );
+ assert!(
+ encoder
+ .decode_value_at(&result[..len], pos, &mut Vec::new())
+ .is_err()
+ );
+ if pos + 1 < row.len() {
+ assert!(
+ encoder
+ .encoded_value_at(&result[..len], pos + 1, &mut Vec::new())
+ .is_err()
+ );
+ }
+ }
+ }
+ assert_eq!(end, result.len());
let mut decoded = Vec::new();
let mut offsets = Vec::new();
// Iter two times to test offsets buffer.
@@ -871,6 +1056,96 @@ mod tests {
}
}
+ #[test]
+ fn test_positional_decode_invalid_encoding() {
+ for (ty, bytes) in [
+ (ConcreteDataType::int64_datatype(), vec![2]),
+ (ConcreteDataType::boolean_datatype(), vec![1, 2]),
+ (ConcreteDataType::binary_datatype(), vec![1, 2]),
+ (ConcreteDataType::binary_datatype(), vec![1, 1, 42, 2]),
+ (ConcreteDataType::string_datatype(), vec![1, 2]),
+ (
+ ConcreteDataType::string_datatype(),
+ vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+ ),
+ (
+ ConcreteDataType::string_datatype(),
+ vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 10],
+ ),
+ ] {
+ let codec = DensePrimaryKeyCodec::with_fields(vec![
+ (0, SortField::new(ty)),
+ (1, SortField::new(ConcreteDataType::int64_datatype())),
+ ]);
+ for pos in [0, 1, 2, usize::MAX] {
+ assert!(
+ codec
+ .encoded_value_at(&bytes, pos, &mut Vec::new())
+ .is_err()
+ );
+ assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err());
+ }
+ assert!(
+ codec
+ .decode_dense_iter(&bytes)
+ .collect::>>()
+ .is_err()
+ );
+ assert!(
+ codec
+ .decode_dense_with(&bytes, &mut Vec::new(), |_, _| {})
+ .is_err()
+ );
+ }
+ let codec = DensePrimaryKeyCodec::with_fields(vec![
+ (0, SortField::new(ConcreteDataType::string_datatype())),
+ (1, SortField::new(ConcreteDataType::int64_datatype())),
+ ]);
+ let bytes = codec
+ .encode([ValueRef::String("abcdefghijk"), ValueRef::Int64(42)].into_iter())
+ .unwrap();
+ for pos in [2, usize::MAX] {
+ assert!(
+ codec
+ .encoded_value_at(&bytes, pos, &mut Vec::new())
+ .is_err()
+ );
+ assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err());
+ }
+ for offset in [bytes.len(), bytes.len() + 1, usize::MAX] {
+ for pos in [0, 1] {
+ assert!(
+ codec
+ .encoded_value_at(&bytes, pos, &mut vec![offset])
+ .is_err()
+ );
+ assert!(
+ codec
+ .decode_value_at(&bytes, pos, &mut vec![offset])
+ .is_err()
+ );
+ }
+ }
+ let mut invalid_utf8 = bytes;
+ invalid_utf8[2] = 0xff;
+ assert!(
+ codec
+ .decode_dense_with(&invalid_utf8, &mut Vec::new(), |_, _| {})
+ .is_err()
+ );
+ assert!(
+ codec
+ .decode_dense_iter(&invalid_utf8)
+ .collect::>>()
+ .is_err()
+ );
+ assert!(
+ codec
+ .decode_value_at(&invalid_utf8, 0, &mut Vec::new())
+ .is_err()
+ );
+ }
+
#[test]
fn test_memcmp_dictionary() {
// Test Dictionary
diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs
index 1e57c13cc7f..0cc9ed8211a 100644
--- a/src/mito-codec/src/row_converter/sparse/checked.rs
+++ b/src/mito-codec/src/row_converter/sparse/checked.rs
@@ -17,6 +17,7 @@ use snafu::{OptionExt, ensure};
use store_api::storage::ColumnId;
use crate::error::{InvalidSparsePrimaryKeySnafu, Result};
+use crate::row_converter::encoded_string_len;
use crate::row_converter::sparse::{
COLUMN_ID_ENCODE_SIZE, RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID,
SparseOffsetsCache, TABLE_ID_VALUE_OFFSET, TAGS_START_OFFSET, TSID_VALUE_OFFSET,
@@ -141,45 +142,13 @@ impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> {
/// Finds the end of an Option without allocating or reading its payload.
/// Unlike memcomparable's skip_bytes, all advances are checked for truncated input.
fn encoded_label(bytes: &[u8]) -> Result<&[u8]> {
- match bytes.first() {
- Some(0) => return Ok(&bytes[..1]),
- Some(1) => {}
- _ => {
- return InvalidSparsePrimaryKeySnafu {
- reason: "invalid label null marker",
- }
- .fail();
+ let len = encoded_string_len(bytes).map_err(|error| {
+ InvalidSparsePrimaryKeySnafu {
+ reason: error.to_string(),
}
- }
- match bytes.get(1) {
- Some(0) => return Ok(&bytes[..2]),
- Some(1) => {}
- _ => {
- return InvalidSparsePrimaryKeySnafu {
- reason: "invalid label bytes marker",
- }
- .fail();
- }
- }
- let mut end = 2;
- loop {
- let chunk = bytes
- .get(end..end + 9)
- .context(InvalidSparsePrimaryKeySnafu {
- reason: "truncated label chunk",
- })?;
- end += 9;
- match chunk[8] {
- 1..=8 => return Ok(&bytes[..end]),
- 9 => {}
- _ => {
- return InvalidSparsePrimaryKeySnafu {
- reason: "invalid label chunk length",
- }
- .fail();
- }
- }
- }
+ .build()
+ })?;
+ Ok(&bytes[..len])
}
#[cfg(test)]
diff --git a/src/mito2/benches/bench_dense_index_update.rs b/src/mito2/benches/bench_dense_index_update.rs
index f81d30eb4a6..6f5ccfc080b 100644
--- a/src/mito2/benches/bench_dense_index_update.rs
+++ b/src/mito2/benches/bench_dense_index_update.rs
@@ -30,8 +30,8 @@ use api::v1::SemanticType;
use common_test_util::temp_dir::create_temp_dir;
use criterion::{Criterion, Throughput, criterion_group, criterion_main};
use datatypes::arrow::array::{
- ArrayRef, BinaryDictionaryBuilder, StringDictionaryBuilder, TimestampMillisecondArray,
- UInt8Array, UInt64Array,
+ ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, StringDictionaryBuilder,
+ TimestampMillisecondArray, UInt8Array, UInt64Array,
};
use datatypes::arrow::datatypes::UInt32Type;
use datatypes::arrow::record_batch::RecordBatch;
@@ -39,6 +39,8 @@ use datatypes::data_type::ConcreteDataType;
use datatypes::schema::{ColumnSchema, SkippingIndexOptions};
use datatypes::value::Value;
use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt};
+use mito2::read::{Batch, BatchBuilder};
+use mito2::sst::index::Indexer;
use mito2::sst::index::intermediate::IntermediateManager;
use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema};
use mito2::test_util::bench_util::{BloomFilterIndexer, InvertedIndexer};
@@ -284,55 +286,115 @@ fn bench_dense_index_update(c: &mut Criterion) {
}
for shape in shapes {
let (metadata, inverted_columns, batch) = input(&shape);
- let run = |iterations| {
- runtime.block_on(async {
- let mut elapsed = Duration::ZERO;
- for _ in 0..iterations {
- let file_id = FileId::random();
- let mut inverted = shape.inverted.then(|| {
- InvertedIndexer::new(
- file_id,
- &metadata,
- intermediate.clone(),
- None,
- NonZeroUsize::new(SEGMENT_ROWS).unwrap(),
- inverted_columns.clone(),
- )
- });
- let mut bloom = if shape.bloom {
- BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None)
- .unwrap()
- } else {
- None
- };
- if allocations {
- ALLOCATIONS.set(Some((0, 0)));
+ for encoded_only in [false, true] {
+ let name = if encoded_only {
+ format!("encoded/{}", shape.name)
+ } else {
+ shape.name.clone()
+ };
+ let encoded_batches: Vec = if encoded_only {
+ let pk = batch
+ .column_by_name("__primary_key")
+ .unwrap()
+ .as_any()
+ .downcast_ref::>()
+ .unwrap();
+ let values = pk.values().as_any().downcast_ref::().unwrap();
+ (0..ROWS)
+ .step_by(shape.rows_per_key)
+ .map(|row| {
+ let mut builder =
+ BatchBuilder::new(values.value(pk.keys().value(row) as usize).to_vec());
+ builder
+ .timestamps_array(
+ batch
+ .column_by_name("ts")
+ .unwrap()
+ .slice(row, shape.rows_per_key),
+ )
+ .unwrap();
+ builder
+ .sequences_array(
+ batch
+ .column_by_name("__sequence")
+ .unwrap()
+ .slice(row, shape.rows_per_key),
+ )
+ .unwrap();
+ builder
+ .op_types_array(
+ batch
+ .column_by_name("__op_type")
+ .unwrap()
+ .slice(row, shape.rows_per_key),
+ )
+ .unwrap();
+ if shape.field {
+ builder
+ .push_field_array(
+ shape.tags,
+ batch
+ .column_by_name(&format!("col_{}", shape.tags))
+ .unwrap()
+ .slice(row, shape.rows_per_key),
+ )
+ .unwrap();
+ }
+ builder.build().unwrap()
+ })
+ .collect()
+ } else {
+ Vec::new()
+ };
+ let run = |iterations| {
+ runtime.block_on(async {
+ let mut elapsed = Duration::ZERO;
+ for _ in 0..iterations {
+ let file_id = FileId::random();
+ let inverted = shape.inverted.then(|| {
+ InvertedIndexer::new(
+ file_id,
+ &metadata,
+ intermediate.clone(),
+ None,
+ NonZeroUsize::new(SEGMENT_ROWS).unwrap(),
+ inverted_columns.clone(),
+ )
+ });
+ let bloom = if shape.bloom {
+ BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None)
+ .unwrap()
+ } else {
+ None
+ };
+ let mut indexer = Indexer::for_bench(&metadata, inverted, bloom);
+ // Fresh caches for each iteration; both creators share each Batch.
+ let mut encoded_batches = encoded_batches.clone();
+ if allocations {
+ ALLOCATIONS.set(Some((0, 0)));
+ }
+ let start = Instant::now();
+ if encoded_only {
+ for batch in &mut encoded_batches {
+ indexer.update(batch).await;
+ }
+ } else {
+ indexer.update_flat(&batch).await;
+ }
+ elapsed += start.elapsed();
+ if allocations {
+ report_allocations(&name);
+ }
+ indexer.abort().await;
}
- let start = Instant::now();
- if let Some(indexer) = &mut inverted {
- indexer.update_flat(&batch).await.unwrap();
- }
- if let Some(indexer) = &mut bloom {
- indexer.update_flat(&batch).await.unwrap();
- }
- elapsed += start.elapsed();
- if allocations {
- report_allocations(&shape.name);
- }
- if let Some(indexer) = &mut inverted {
- indexer.abort().await.unwrap();
- }
- if let Some(indexer) = &mut bloom {
- indexer.abort().await.unwrap();
- }
- }
- elapsed
- })
- };
- if allocations {
- run(1);
- } else {
- group.bench_function(&shape.name, |b| b.iter_custom(&run));
+ elapsed
+ })
+ };
+ if allocations {
+ run(1);
+ } else {
+ group.bench_function(&name, |b| b.iter_custom(&run));
+ }
}
}
group.finish();
diff --git a/src/mito2/benches/bench_pk_tag_column.rs b/src/mito2/benches/bench_pk_tag_column.rs
index d3ded3d0b05..9e191a69511 100644
--- a/src/mito2/benches/bench_pk_tag_column.rs
+++ b/src/mito2/benches/bench_pk_tag_column.rs
@@ -24,6 +24,7 @@
use std::hint::black_box;
use std::sync::Arc;
+use std::time::Duration;
use api::v1::SemanticType;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
@@ -35,9 +36,12 @@ use datatypes::arrow::datatypes::UInt32Type;
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::ColumnSchema;
-use mito_codec::row_converter::SparsePrimaryKeyCodec;
+use datatypes::value::Value;
+use mito_codec::row_converter::{
+ DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField, SparsePrimaryKeyCodec,
+};
use mito2::sst::parquet::flat_format::decode_primary_keys;
-use mito2::test_util::bench_util::tag_filter_for_bench;
+use mito2::test_util::bench_util::{pk_materializer_for_bench, tag_filter_for_bench};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef};
use store_api::storage::RegionId;
@@ -237,5 +241,112 @@ fn bench_pk_tag_filters(c: &mut Criterion) {
group.finish();
}
-criterion_group!(benches, bench_pk_tag_column, bench_pk_tag_filters);
+/// Measures encoded-only Dense inputs, including the cost of discovering offsets.
+fn bench_dense_pk_tag_column(c: &mut Criterion) {
+ const TAGS: usize = 40;
+ let mut group = c.benchmark_group("dense_pk_tag_column");
+ group.sample_size(30);
+ group.warm_up_time(Duration::from_secs(1));
+ group.measurement_time(Duration::from_secs(3));
+ for numeric in [false, true] {
+ let ty = if numeric {
+ ConcreteDataType::uint64_datatype()
+ } else {
+ ConcreteDataType::string_datatype()
+ };
+ let codec = DensePrimaryKeyCodec::with_fields(
+ (0..TAGS)
+ .map(|id| (id as u32, SortField::new(ty.clone())))
+ .collect(),
+ );
+ let mut metadata = RegionMetadataBuilder::new(RegionId::new(1, 1));
+ for id in 0..TAGS {
+ metadata.push_column_metadata(ColumnMetadata {
+ column_id: id as u32,
+ column_schema: ColumnSchema::new(format!("tag_{id}"), ty.clone(), true),
+ semantic_type: SemanticType::Tag,
+ });
+ }
+ metadata
+ .push_column_metadata(ColumnMetadata {
+ column_id: TAGS as u32,
+ column_schema: ColumnSchema::new(
+ "ts",
+ ConcreteDataType::timestamp_millisecond_datatype(),
+ false,
+ ),
+ semantic_type: SemanticType::Timestamp,
+ })
+ .primary_key((0..TAGS as u32).collect());
+ let metadata = Arc::new(metadata.build().unwrap());
+ for rows_per_key in [1, 32] {
+ let mut keys = BinaryDictionaryBuilder::::new();
+ for series in 0..ROWS / rows_per_key {
+ let values: Vec<_> = (0..TAGS)
+ .map(|id| {
+ if numeric {
+ Value::UInt64((series + id) as u64)
+ } else {
+ Value::from(format!(
+ "tag-{id:03}-value-{series:010}-abcdefghijklmnopqrstuvwxyz"
+ ))
+ }
+ })
+ .collect();
+ let pk = codec
+ .encode(values.iter().map(Value::as_value_ref))
+ .unwrap();
+ for _ in 0..rows_per_key {
+ keys.append(&pk).unwrap();
+ }
+ }
+ let batch = RecordBatch::try_from_iter([
+ (
+ "ts",
+ Arc::new(TimestampMillisecondArray::from_iter_values(0..ROWS as i64))
+ as ArrayRef,
+ ),
+ ("__primary_key", Arc::new(keys.finish()) as ArrayRef),
+ (
+ "__sequence",
+ Arc::new(UInt64Array::from(vec![1; ROWS])) as ArrayRef,
+ ),
+ (
+ "__op_type",
+ Arc::new(UInt8Array::from(vec![1; ROWS])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ for (projection, start, count) in [
+ ("first", 0, 1),
+ ("last", 39, 1),
+ ("4last", 36, 4),
+ ("all", 0, 40),
+ ] {
+ let kind = if numeric { "numeric" } else { "string" };
+ let materialize = pk_materializer_for_bench(
+ metadata.clone(),
+ (start as u32..(start + count) as u32)
+ .chain([TAGS as u32])
+ .collect(),
+ batch.schema(),
+ );
+ assert_eq!(materialize(batch.clone()).unwrap().num_columns(), count + 4);
+ group.bench_function(format!("{kind}_{projection}_{rows_per_key}rpk"), |b| {
+ b.iter(|| {
+ black_box(materialize(black_box(batch.clone())).unwrap());
+ });
+ });
+ }
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(
+ benches,
+ bench_pk_tag_column,
+ bench_pk_tag_filters,
+ bench_dense_pk_tag_column
+);
criterion_main!(benches);
diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs
index 337386aa2ff..685cc0d2fc7 100644
--- a/src/mito2/src/access_layer.rs
+++ b/src/mito2/src/access_layer.rs
@@ -18,6 +18,7 @@ use std::time::{Duration, Instant};
use async_stream::try_stream;
use common_base::readable_size::ReadableSize;
+use common_runtime::runtime::RuntimeTrait;
use common_telemetry::warn;
use common_time::Timestamp;
use futures::{Stream, TryStreamExt};
@@ -348,6 +349,7 @@ impl AccessLayer {
write_opts: &WriteOptions,
metrics: &mut Metrics,
) -> Result {
+ let op_type = request.op_type;
let region_id = request.metadata.region_id;
let region_metadata = request.metadata.clone();
let cache_manager = request.cache_manager.clone();
@@ -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),
diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs
index 0d43d31f434..ef150dfb070 100644
--- a/src/mito2/src/cache.rs
+++ b/src/mito2/src/cache.rs
@@ -31,6 +31,8 @@ use std::sync::{Arc, RwLock, Weak};
use bytes::Bytes;
use common_base::readable_size::ReadableSize;
use common_datasource::compression::CompressionType;
+use common_runtime::Runtime;
+use common_runtime::runtime::RuntimeTrait;
use common_telemetry::warn;
use datatypes::arrow::buffer::BooleanBuffer;
use datatypes::arrow::record_batch::RecordBatch;
@@ -238,7 +240,7 @@ impl SstMetaPreparation {
}
impl CompactSstMeta {
- async fn decode(&self) -> Result> {
+ async fn decode(&self, runtime: &Runtime) -> Result> {
let mut decoded_guard = self.decoded.lock().await;
if let Some(decoded) = decoded_guard.upgrade() {
return Ok(decoded);
@@ -248,33 +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,
page_index_policy: PageIndexPolicy,
+ runtime: &Runtime,
) -> Result> {
let file_path = file_path.to_string();
- common_runtime::spawn_blocking_global(move || {
- let parquet_metadata = strip_column_indexes(parquet_metadata);
- CachedSstMeta::try_new_with_page_index_policy(
- &file_path,
- parquet_metadata,
- region_metadata,
- page_index_policy,
- )
- .map(Arc::new)
- })
- .await
- .context(JoinSnafu)?
+ runtime
+ .spawn_blocking(move || {
+ let parquet_metadata = strip_column_indexes(parquet_metadata);
+ CachedSstMeta::try_new_with_page_index_policy(
+ &file_path,
+ parquet_metadata,
+ region_metadata,
+ page_index_policy,
+ )
+ .map(Arc::new)
+ })
+ .await
+ .context(JoinSnafu)?
}
-/// Decodes SST metadata and attempts to encode both cache representations on the blocking runtime.
+/// Decodes SST metadata and attempts to encode both cache representations on the given blocking runtime.
pub(crate) async fn prepare_sst_meta(
file_path: &str,
parquet_metadata: ParquetMetaData,
region_metadata: Option,
page_index_policy: PageIndexPolicy,
+ runtime: &Runtime,
) -> Result {
let file_path = file_path.to_string();
- common_runtime::spawn_blocking_global(move || {
- prepare_sst_meta_sync(
- &file_path,
- parquet_metadata,
- region_metadata,
- page_index_policy,
- )
- })
- .await
- .context(JoinSnafu)?
+ runtime
+ .spawn_blocking(move || {
+ prepare_sst_meta_sync(
+ &file_path,
+ parquet_metadata,
+ region_metadata,
+ page_index_policy,
+ )
+ })
+ .await
+ .context(JoinSnafu)?
}
/// Synchronously prepares SST metadata. Callers must run this on a blocking runtime.
@@ -671,6 +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> {
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(®ion_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()
);
diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs
index 0594a1ac10d..4cb60ed8d48 100644
--- a/src/mito2/src/cache/file_cache.rs
+++ b/src/mito2/src/cache/file_cache.rs
@@ -21,6 +21,7 @@ use std::time::{Duration, Instant};
use bytes::Bytes;
use common_base::readable_size::ReadableSize;
+use common_runtime::Runtime;
use common_telemetry::{debug, error, info, warn};
use futures::{AsyncWriteExt, FutureExt, TryStreamExt};
use moka::future::Cache;
@@ -627,12 +628,13 @@ impl FileCache {
key: IndexKey,
cache_metrics: &mut MetadataCacheMetrics,
page_index_policy: PageIndexPolicy,
+ runtime: &Runtime,
) -> Option {
let file_path = self.inner.cache_file_path(key);
let metadata = self
.get_parquet_meta_data(key, cache_metrics, page_index_policy)
.await?;
- match prepare_sst_meta(&file_path, metadata, None, page_index_policy).await {
+ match prepare_sst_meta(&file_path, metadata, None, page_index_policy, runtime).await {
Ok(metadata) => Some(metadata),
Err(err) => {
CACHE_MISS
@@ -653,12 +655,13 @@ impl FileCache {
key: IndexKey,
cache_metrics: &mut MetadataCacheMetrics,
page_index_policy: PageIndexPolicy,
+ runtime: &Runtime,
) -> Option> {
let file_path = self.inner.cache_file_path(key);
let metadata = self
.get_parquet_meta_data(key, cache_metrics, page_index_policy)
.await?;
- match decode_sst_meta(&file_path, metadata, None, page_index_policy).await {
+ match decode_sst_meta(&file_path, metadata, None, page_index_policy, runtime).await {
Ok(metadata) => Some(metadata),
Err(err) => {
CACHE_MISS
diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs
index f381f80722e..cb2a9d34ec8 100644
--- a/src/mito2/src/compaction/json2.rs
+++ b/src/mito2/src/compaction/json2.rs
@@ -19,7 +19,9 @@ use arrow_schema::extension::ExtensionType;
use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, SchemaRef};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::extension::json::{JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, JsonMetadata};
-use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint};
+use datatypes::json::{
+ JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint, TypeHintMismatchPolicy,
+};
use datatypes::prelude::ConcreteDataType;
use datatypes::types::json_type::JsonNativeType;
use datatypes::vectors::json::array::JsonArray;
@@ -260,7 +262,7 @@ fn select_dynamic_hints(
.filter(|(path, stat)| {
!stat.is_type_conflicted
// TODO(LFC): Instead of "primitive only", consider retaining stable compound types
- // that are safe to write to Parquet, as flush does. Or better, unite the two
+ // that are safe to write to Parquet, as flush does. Or better, unite the two
// selection process.
&& stat.data_type.is_primitive()
&& !has_ancestor_path(path)
@@ -327,7 +329,12 @@ pub(crate) fn rewrite_json2_batch(
};
let array = JsonArray::from(array)
- .rewrite_to_v2(field, &plan.logical_settings, &plan.target_layout)
+ .rewrite_to_v2_with_type_hint_mismatch_policy(
+ field,
+ &plan.logical_settings,
+ &plan.target_layout,
+ TypeHintMismatchPolicy::CoerceOrNull,
+ )
.context(ConvertValueSnafu)?;
debug_assert_eq!(
&json2_physical_data_type(&plan.target_layout),
diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs
index 72e778851ee..d5abe29f021 100644
--- a/src/mito2/src/read.rs
+++ b/src/mito2/src/read.rs
@@ -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,
/// Possibly decoded `primary_key` values. Some places would decode it in advance.
pk_values: Option,
+ /// Lazily decoded Dense columns, shared by consumers such as index creators.
+ dense_pk_cache: Option>,
/// Timestamps of rows, should be sorted and not null.
timestamps: VectorRef,
/// Sequences of rows
@@ -135,6 +137,13 @@ pub struct Batch {
fields_idx: Option>,
}
+/// Only used when no fully decoded primary key was supplied by the reader.
+#[derive(Debug, PartialEq, Clone)]
+struct DensePkCache {
+ offsets: Vec,
+ values: Vec>,
+}
+
impl Batch {
/// Creates a new batch.
pub fn new(
@@ -173,12 +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) {
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> {
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],
diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs
index 19c487ed0b8..ef09f5f48ab 100644
--- a/src/mito2/src/read/pruner.rs
+++ b/src/mito2/src/read/pruner.rs
@@ -350,12 +350,14 @@ impl Pruner {
enable_predicate_prefilter,
});
- // Spawn worker tasks with their receivers
+ // Keep pruning and prefetching on the runtime of the originating workload.
for (worker_id, rx) in receivers.into_iter().enumerate() {
- let inner_clone = inner.clone();
- common_runtime::spawn_query(async move {
- Self::worker_loop(worker_id, rx, inner_clone).await;
- });
+ let worker = Self::worker_loop(worker_id, rx, inner.clone());
+ if inner.stream_ctx.input.compaction {
+ common_runtime::spawn_compact(worker);
+ } else {
+ common_runtime::spawn_query(worker);
+ }
}
Self {
diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs
index fb5d260fe6f..84a0e13833f 100644
--- a/src/mito2/src/region/opener.rs
+++ b/src/mito2/src/region/opener.rs
@@ -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
{
diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs
index 0c3ede4fb06..2dd29e9803e 100644
--- a/src/mito2/src/sst/index.rs
+++ b/src/mito2/src/sst/index.rs
@@ -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,
last_mem_bloom_filter: usize,
+ /// Present only when the active creators together need every Dense PK field.
+ dense_pk_decoder: Option,
#[cfg(feature = "vector_index")]
vector_indexer: Option,
#[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,
+ bloom_filter_indexer: Option,
+ ) -> Self {
+ let mut indexer = Self {
+ region_id: metadata.region_id,
+ inverted_indexer,
+ bloom_filter_indexer,
+ ..Default::default()
+ };
+ indexer.prepare_dense_pk_decoder(metadata);
+ indexer
+ }
+
+ /// Computes demand once from active creators. Field columns and duplicate
+ /// requests across creators cannot turn a partial PK projection into a full one.
+ fn prepare_dense_pk_decoder(&mut self, metadata: &RegionMetadataRef) {
+ if metadata.primary_key_encoding != PrimaryKeyEncoding::Dense
+ || metadata.primary_key.is_empty()
+ {
+ return;
+ }
+ let requested: HashSet<_> = self
+ .inverted_indexer
+ .iter()
+ .flat_map(|indexer| indexer.column_ids())
+ .chain(
+ self.bloom_filter_indexer
+ .iter()
+ .flat_map(|indexer| indexer.column_ids()),
+ )
+ .collect();
+ if metadata.primary_key.iter().all(|id| requested.contains(id)) {
+ self.dense_pk_decoder = Some(DensePrimaryKeyCodec::new(metadata));
+ }
+ }
+
+ /// Called before any creator consumes the batch, so the full decode is shared.
+ fn prepare_primary_key(&self, batch: &mut Batch) -> Result<()> {
+ if let Some(codec) = &self.dense_pk_decoder {
+ batch.ensure_dense_pk_decoded(codec)?;
+ }
+ Ok(())
+ }
+
/// Updates the index with the given batch.
pub async fn update(&mut self, batch: &mut Batch) {
self.do_update(batch).await;
@@ -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!(
diff --git a/src/mito2/src/sst/index/column_test.rs b/src/mito2/src/sst/index/column_test.rs
index 146bed120c0..34ff1c19fde 100644
--- a/src/mito2/src/sst/index/column_test.rs
+++ b/src/mito2/src/sst/index/column_test.rs
@@ -29,7 +29,9 @@ use index::bitmap::{Bitmap, BitmapType};
use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl};
use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader};
use mito_codec::index::IndexValueCodec;
-use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField};
+use mito_codec::row_converter::{
+ DensePrimaryKeyCodec, PrimaryKeyCodec, PrimaryKeyCodecExt, SortField,
+};
use object_store::ObjectStore;
use object_store::services::Memory;
use prost::Message;
@@ -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" => {
diff --git a/src/mito2/src/sst/index/indexer/abort.rs b/src/mito2/src/sst/index/indexer/abort.rs
index a928eb186da..24cbb797cbd 100644
--- a/src/mito2/src/sst/index/indexer/abort.rs
+++ b/src/mito2/src/sst/index/indexer/abort.rs
@@ -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;
diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs
index c7966823ddf..823cd5c2b8a 100644
--- a/src/mito2/src/sst/index/indexer/finish.rs
+++ b/src/mito2/src/sst/index/indexer/finish.rs
@@ -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 {
diff --git a/src/mito2/src/sst/index/indexer/update.rs b/src/mito2/src/sst/index/indexer/update.rs
index 095da337f49..0168bee876e 100644
--- a/src/mito2/src/sst/index/indexer/update.rs
+++ b/src/mito2/src/sst/index/indexer/update.rs
@@ -24,6 +24,11 @@ impl Indexer {
return;
}
+ if !self.do_prepare_primary_key(batch) {
+ self.do_abort().await;
+ return;
+ }
+
if !self.do_update_inverted_index(batch).await {
self.do_abort().await;
}
@@ -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 {
diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs
index a99fc201baf..34fb8cb873c 100644
--- a/src/mito2/src/sst/parquet/flat_format.rs
+++ b/src/mito2/src/sst/parquet/flat_format.rs
@@ -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 {
- 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),
+ /// Dense keys share positional offsets across projected tag columns.
+ Dense {
+ codec: DensePrimaryKeyCodec,
+ values: BinaryArray,
+ distinct_keys: Vec,
+ /// Initialized by positional access; full extraction needs no offsets.
+ offsets: Vec>,
+ },
/// Sparse primary keys stay encoded; each tag column extracts only its own
/// values from the raw keys.
Sparse {
@@ -811,7 +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> {
+ 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::::new(
+ UInt32Array::from(row_keys.to_vec()),
+ Arc::new(BinaryArray::from_iter_values(&encoded)),
+ );
+ let batch = RecordBatch::try_from_iter([
+ (
+ "ts",
+ Arc::new(TimestampMillisecondArray::from_iter_values(0..6)) as ArrayRef,
+ ),
+ ("__primary_key", Arc::new(pk) as ArrayRef),
+ (
+ "__sequence",
+ Arc::new(UInt64Array::from(vec![1; 6])) as ArrayRef,
+ ),
+ (
+ "__op_type",
+ Arc::new(UInt8Array::from(vec![1; 6])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ let mut metadata = RegionMetadataBuilder::new(RegionId::new(1, 1));
+ for (id, ty) in columns.iter().rev() {
+ metadata.push_column_metadata(ColumnMetadata {
+ column_id: *id,
+ column_schema: ColumnSchema::new(format!("tag_{id}"), ty.clone(), true),
+ semantic_type: SemanticType::Tag,
+ });
+ }
+ metadata
+ .push_column_metadata(ColumnMetadata {
+ column_id: 23,
+ column_schema: ColumnSchema::new(
+ "ts",
+ ConcreteDataType::timestamp_millisecond_datatype(),
+ false,
+ ),
+ semantic_type: SemanticType::Timestamp,
+ })
+ .primary_key(vec![17, 9, 2]);
+ let metadata = Arc::new(metadata.build().unwrap());
+ let format = FlatReadFormat::new(
+ metadata.clone(),
+ ReadColumns::new([17, 9, 2, 23]),
+ Some(batch.schema()),
+ "test",
+ false,
+ )
+ .unwrap();
+ for (offset, len) in [(0, 6), (1, 4), (3, 0)] {
+ let mut decoded = decode_primary_keys(&codec, &batch.slice(offset, len)).unwrap();
+ let all = decoded.get_dense_tag_columns().unwrap();
+ let converted = format
+ .convert_batch(batch.slice(offset, len), None)
+ .unwrap();
+ // Out-of-order and repeated projections share the same key offsets.
+ for pos in [2, 0, 1, 0] {
+ let array = decoded
+ .get_tag_column(columns[pos].0, Some(pos), &columns[pos].1)
+ .unwrap();
+ assert_eq!(array.to_data(), all[pos].to_data());
+ assert_eq!(converted.column(pos).to_data(), all[pos].to_data());
+ let actual = Helper::try_into_vector(array).unwrap();
+ for (row, &key) in row_keys[offset..offset + len].iter().enumerate() {
+ let eager = codec
+ .decode_dense_without_column_id(&encoded[key as usize])
+ .unwrap();
+ assert_eq!(actual.get(row), eager[pos]);
+ }
+ }
+ // A position beyond the source schema remains NULL, never decoded
+ // with a newer schema's field order or type.
+ let missing = decoded
+ .get_tag_column(99, Some(3), &ConcreteDataType::int64_datatype())
+ .unwrap();
+ assert_eq!(missing.null_count(), len);
+ assert!(decoded.get_tag_column(17, None, &columns[0].1).is_err());
+ }
+ let mut malformed_columns = batch.columns().to_vec();
+ malformed_columns[1] = Arc::new(DictionaryArray::::new(
+ UInt32Array::from(vec![0; 6]),
+ Arc::new(BinaryArray::from(vec![&[0, 1][..]])),
+ ));
+ let malformed = RecordBatch::try_new(batch.schema(), malformed_columns).unwrap();
+ assert!(format.convert_batch(malformed.clone(), None).is_err());
+ // A partial projection must still avoid decoding an unneeded broken suffix.
+ let partial = FlatReadFormat::new(
+ metadata,
+ ReadColumns::new([17, 23]),
+ Some(batch.schema()),
+ "test",
+ false,
+ )
+ .unwrap();
+ let partial = partial.convert_batch(malformed, None).unwrap();
+ let tag = Helper::try_into_vector(partial.column(0).clone()).unwrap();
+ assert!((0..6).all(|row| tag.get(row).is_null()));
+ }
+
/// Builds a `RegionMetadata` with the given number of tags and fields.
fn build_metadata(
num_tags: usize,
diff --git a/src/mito2/src/sst/parquet/json_align.rs b/src/mito2/src/sst/parquet/json_align.rs
index 8af88f6ac0f..93ecffa072f 100644
--- a/src/mito2/src/sst/parquet/json_align.rs
+++ b/src/mito2/src/sst/parquet/json_align.rs
@@ -22,7 +22,7 @@ use datatypes::arrow::array::{ArrayRef, new_null_array};
use datatypes::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::extension::json::{JsonMetadata, is_json2_extension_type};
-use datatypes::json::JsonSettings;
+use datatypes::json::{JsonSettings, TypeHintMismatchPolicy};
use datatypes::vectors::json::array::JsonArray;
use datatypes::vectors::json::json2_physical_data_type;
use futures::Stream;
@@ -251,10 +251,11 @@ fn rewrite_array(
settings: &RewriteSettings,
) -> Result {
JsonArray::from(source_array)
- .rewrite_to_v2(
+ .rewrite_to_v2_with_type_hint_mismatch_policy(
source_field,
&settings.logical_settings,
&settings.target_layout,
+ TypeHintMismatchPolicy::CoerceOrNull,
)
.context(DataTypeMismatchSnafu)
}
@@ -330,7 +331,9 @@ mod tests {
Array, ArrayRef, BinaryArray, Int64Array, StringArray, StringViewArray, StructArray,
};
use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema};
- use datatypes::extension::json::Json2ExtensionType;
+ use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
+ use datatypes::json::JsonTypeHint;
+ use datatypes::prelude::ConcreteDataType;
use datatypes::types::parse_string_to_jsonb;
use futures::{StreamExt, stream};
@@ -783,6 +786,75 @@ mod tests {
assert_eq!(int_array([10, 20]).as_ref(), output.column(3).as_ref());
}
+ #[tokio::test]
+ async fn test_rewrite_keeps_rows_with_invalid_json2_settings() {
+ let settings = JsonSettings::try_new(
+ vec![JsonTypeHint {
+ path: vec!["kind".to_string()],
+ data_type: ConcreteDataType::string_datatype(),
+ inverted_index: false,
+ }],
+ Some(0),
+ )
+ .unwrap();
+ let target_type = json2_physical_data_type(&settings);
+ let output_schema = schema([
+ Field::new("j", target_type.clone(), true).with_extension_type(
+ Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings.clone()))),
+ ),
+ Field::new("value", DataType::Int64, true),
+ ]);
+ let source = Arc::new(BinaryArray::from_iter([
+ Some(parse_string_to_jsonb(r#"{"kind":"valid"}"#).unwrap()),
+ Some(parse_string_to_jsonb(r#"{"kind":1}"#).unwrap()),
+ ])) as ArrayRef;
+ let input = RecordBatch::try_new(
+ schema([
+ Field::new("j", DataType::Binary, true)
+ .with_extension_type(Json2ExtensionType::default()),
+ Field::new("value", DataType::Int64, true),
+ ]),
+ vec![source, int_array([10, 20])],
+ )
+ .unwrap();
+ let columns = HashMap::from([(
+ "j".to_string(),
+ Json2TargetLayout {
+ extension_metadata: serde_json::to_string(&JsonMetadata::new(settings.clone()))
+ .unwrap(),
+ target_layout: settings,
+ },
+ )]);
+ let mut aligner = JsonSchemaAligner::new(
+ stream::iter([Ok(input)]),
+ vec![true, true],
+ output_schema,
+ AlignMode::Rewrite { columns },
+ )
+ .unwrap();
+
+ let output = aligner.next().await.unwrap().unwrap();
+ assert_eq!(2, output.num_rows());
+ assert_eq!(
+ 10,
+ output
+ .column(1)
+ .as_any()
+ .downcast_ref::()
+ .unwrap()
+ .value(0)
+ );
+ assert_eq!(
+ 20,
+ output
+ .column(1)
+ .as_any()
+ .downcast_ref::()
+ .unwrap()
+ .value(1)
+ );
+ }
+
#[test]
fn test_rewrite_rejects_mismatched_output_layout() {
let columns = HashMap::from([(
diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs
index f54812f97d4..9584ff0a6af 100644
--- a/src/mito2/src/sst/parquet/reader.rs
+++ b/src/mito2/src/sst/parquet/reader.rs
@@ -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) => {
diff --git a/src/mito2/src/test_util/bench_util.rs b/src/mito2/src/test_util/bench_util.rs
index ac635d2799a..e0334676e5b 100644
--- a/src/mito2/src/test_util/bench_util.rs
+++ b/src/mito2/src/test_util/bench_util.rs
@@ -21,6 +21,7 @@ use api::v1::value::ValueData;
use api::v1::{Row, Rows, SemanticType};
use datafusion_common::Column;
use datafusion_expr::{Expr, lit};
+use datatypes::arrow::datatypes::SchemaRef;
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::ColumnSchema;
@@ -30,7 +31,7 @@ use rand::seq::IndexedRandom;
use store_api::metadata::{
ColumnMetadata, RegionMetadata, RegionMetadataBuilder, RegionMetadataRef,
};
-use store_api::storage::RegionId;
+use store_api::storage::{ColumnId, RegionId};
use table::predicate::Predicate;
use crate::memtable::KeyValues;
@@ -42,6 +43,23 @@ use crate::sst::parquet::flat_format::FlatReadFormat;
use crate::sst::parquet::reader::SimpleFilterContext;
use crate::test_util::memtable_util::region_metadata_to_row_schema;
+/// Builds the actual encoded-PK-to-flat conversion used by the SST reader.
+pub fn pk_materializer_for_bench(
+ metadata: RegionMetadataRef,
+ columns: Vec,
+ file_schema: SchemaRef,
+) -> impl Fn(RecordBatch) -> crate::error::Result {
+ let format = FlatReadFormat::new(
+ metadata,
+ ReadColumns::new(columns),
+ Some(file_schema),
+ "bench",
+ false,
+ )
+ .unwrap();
+ move |batch| format.convert_batch(batch, None)
+}
+
/// Builds a precise-filter benchmark with tags left encoded in the primary key.
pub fn tag_filter_for_bench(
metadata: RegionMetadataRef,
diff --git a/src/object-store/src/test_util.rs b/src/object-store/src/test_util.rs
index af68978b0e4..936ecba5bac 100644
--- a/src/object-store/src/test_util.rs
+++ b/src/object-store/src/test_util.rs
@@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-use std::env;
-
use crate::{ObjectStore, Result};
/// Temp folder for object store test
@@ -35,30 +33,3 @@ impl TempFolder {
self.store.delete_with(&self.path).recursive(true).await
}
}
-
-/// Test s3 config from environment variables
-#[derive(Debug)]
-pub struct TestS3Config {
- pub root: String,
- pub access_key_id: String,
- pub secret_access_key: String,
- pub bucket: String,
- pub region: Option,
-}
-
-/// Returns s3 test config, return None if not found.
-pub fn s3_test_config() -> Option {
- if let Ok(b) = env::var("GT_S3_BUCKET")
- && !b.is_empty()
- {
- return Some(TestS3Config {
- root: uuid::Uuid::new_v4().to_string(),
- access_key_id: env::var("GT_S3_ACCESS_KEY_ID").ok()?,
- secret_access_key: env::var("GT_S3_ACCESS_KEY").ok()?,
- bucket: env::var("GT_S3_BUCKET").ok()?,
- region: Some(env::var("GT_S3_REGION").ok()?),
- });
- }
-
- None
-}
diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs
index 9cf03fe42c5..735d2c56f15 100644
--- a/src/object-store/tests/object_store_test.rs
+++ b/src/object-store/tests/object_store_test.rs
@@ -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);
diff --git a/src/operator/src/batcher.rs b/src/operator/src/batcher.rs
index 4c5e9892c0b..677524d53f3 100644
--- a/src/operator/src/batcher.rs
+++ b/src/operator/src/batcher.rs
@@ -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>;
- /// 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,
diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs
index b28af78d383..1d06c24ff7f 100644
--- a/src/operator/src/expr_helper.rs
+++ b/src/operator/src/expr_helper.rs
@@ -24,9 +24,10 @@ use api::v1::column_def::{options_from_column_schema, try_as_column_schema};
use api::v1::{
AddColumn, AddColumns, AlterDatabaseExpr, AlterTableExpr, Analyzer, ColumnDataType,
ColumnDataTypeExtension, CreateFlowExpr, CreateTableExpr, CreateViewExpr, DropColumn,
- DropColumns, DropDefaults, ExpireAfter, FulltextBackend as PbFulltextBackend, ModifyColumnType,
+ DropColumns, DropDefaults, ExpireAfter, FulltextBackend as PbFulltextBackend,
+ JsonSettings as PbJsonSettings, JsonTypeHint as PbJsonTypeHint, ModifyColumnType,
ModifyColumnTypes, RenameTable, SemanticType, SetDatabaseOptions, SetDefaults, SetFulltext,
- SetIndex, SetIndexes, SetInverted, SetSkipping, SetTableOptions,
+ SetIndex, SetIndexes, SetInverted, SetJsonSettings, SetSkipping, SetTableOptions,
SkippingIndexType as PbSkippingIndexType, TableName, UnsetDatabaseOptions, UnsetFulltext,
UnsetIndex, UnsetIndexes, UnsetInverted, UnsetSkipping, UnsetTableOptions, set_index,
unset_index,
@@ -36,6 +37,7 @@ use common_error::ext::BoxedError;
use common_grpc_expr::util::ColumnExpr;
use common_time::Timezone;
use datafusion::sql::planner::object_name_to_table_reference;
+use datatypes::json::JsonSettings;
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::{
COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND,
@@ -782,6 +784,29 @@ pub(crate) fn to_repartition_request(
})
}
+fn json_settings_to_proto(settings: JsonSettings) -> Result {
+ let (type_hints, max_auto_expanded_paths) = settings.into_parts();
+ let type_hints = type_hints
+ .into_iter()
+ .map(|hint| {
+ let (data_type, datatype_extension) = ColumnDataTypeWrapper::try_from(hint.data_type)
+ .map(|w| w.to_parts())
+ .context(ColumnDataTypeSnafu)?;
+
+ Ok(PbJsonTypeHint {
+ path: hint.path,
+ data_type: data_type as i32,
+ datatype_extension,
+ })
+ })
+ .collect::>>()?;
+
+ Ok(PbJsonSettings {
+ type_hints,
+ max_auto_expanded_paths,
+ })
+}
+
/// Converts a SQL alter table statement into a gRPC alter table expression.
pub(crate) fn to_alter_table_expr(
alter_table: AlterTable,
@@ -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]
diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs
index afc903621e2..49215f89713 100644
--- a/src/operator/src/insert.rs
+++ b/src/operator/src/insert.rs
@@ -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 {
+ 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 {
+ 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::>();
+ assert_eq!(
+ totals,
+ if enabled {
+ vec![("a", 4, 0), ("b", 2, 0)]
+ } else {
+ vec![]
+ }
+ );
+ }
+ meter.attempts.lock().unwrap().clear();
+ meter.reject.store(false, Ordering::Relaxed);
+ let ctx = Arc::new(QueryContext::with_channel(
+ CATALOG,
+ "logical",
+ Channel::Otlp,
+ ));
+ let admitted = admit_write(2, &ctx).await.unwrap();
+ let mut requests = RowInsertRequests {
+ inserts: vec![RowInsertRequest {
+ table_name: "metric".to_string(),
+ rows: Some(Rows {
+ schema: vec![],
+ rows: vec![api::v1::Row::default(); 2],
+ }),
+ }],
+ };
+ let original = requests.clone();
+ let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
+ .await
+ .unwrap();
+ assert_eq!(cost, if enabled { 17 } else { 0 });
+ assert_eq!(requests, original);
+ let records = meter
+ .attempts
+ .lock()
+ .unwrap()
.iter()
- .map(|r| (r.schema.as_str(), r.rows, r.value))
+ .map(|record| (record.rows, record.value))
.collect::>();
assert_eq!(
- totals,
+ records,
if enabled {
- vec![("a", 4, 0), ("b", 2, 0)]
+ vec![(2, 0), (0, 17)]
} else {
vec![]
}
);
+ meter.reject.store(true, Ordering::Relaxed);
+ let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
+ if enabled {
+ assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
+ } else {
+ assert_eq!(result.unwrap(), 0);
+ }
+ assert_eq!(requests, original);
}
#[test]
@@ -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 = Arc::new(UnexpectedBatcher);
diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs
index 94081c2d284..44211b8d2b7 100644
--- a/src/query/src/promql/planner.rs
+++ b/src/query/src/promql/planner.rs
@@ -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]
diff --git a/src/query/src/promql/planner/matching_filters.rs b/src/query/src/promql/planner/matching_filters.rs
index 386d95e9228..e38d48044da 100644
--- a/src/query/src/promql/planner/matching_filters.rs
+++ b/src/query/src/promql/planner/matching_filters.rs
@@ -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 {
+ 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]) -> Option {
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])"#,
diff --git a/src/servers/src/batcher.rs b/src/servers/src/batcher.rs
index 10f8f558eb3..6b6ca8f7ee7 100644
--- a/src/servers/src/batcher.rs
+++ b/src/servers/src/batcher.rs
@@ -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::().ok())
+ .unwrap_or(true)
+}
diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs
index f4c764ee8a8..bd7a26da613 100644
--- a/src/servers/src/batcher/logical_table.rs
+++ b/src/servers/src/batcher/logical_table.rs
@@ -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::().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 {
+ 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(
+ &self,
+ requests: RowInsertRequests,
+ ctx: QueryContextRef,
+ after_prepare: impl FnOnce(RowInsertRequests) -> F,
+ ) -> Result<(u64, T)>
+ where
+ F: Future